[发明专利]图像处理模型的训练方法及装置、计算机可读存储介质

说明书

本申请实施例提供一种图像处理模型的训练方法及装置、计算机可读存储介质，所述图像处理模型包括U型网络，所述U型网络包括下采样通路和上采样通路，所述下采样通路包括第一下采样模块，所述上采样通道包括第一上采样模块，所述第一上采样模块与所述第一下采样模块相连，以复用所述第一下采用模块输出的特征，所述方法包括：将训练数据输入所述图像处理模型，得到所述图像处理模型的训练损失；根据所述训练损失，更新所述图像处理模型的模型参数；在更新所述模型参数的过程中，对所述第一下采样模块和所述第一上采样模块之间的网络架构进行神经网络架构搜索，以确定所述第一下采样模块和所述第一上采样模块之间的注意力模块的结构。

技术领域

本申请涉及图像处理技术领域，具体涉及一种图像处理模型的训练方法及装置、计算机可读存储介质。

背景技术

在基于U型网的图像处理模型中，U型网的下采样模块和上采样模块之间的注意力结构通常由人工设计而成，存在图形处理性能较差的问题。

发明内容

为了解决上述问题，本申请实施例提供一种图像处理模型的训练方法及装置、计算机可读存储介质。

第一方面，提供一种图像处理模型的训练方法，所述图像处理模型包括U型网络，所述U型网络包括下采样通路和上采样通路，所述下采样通路包括第一下采样模块，所述上采样通道包括第一上采样模块，所述第一上采样模块与所述第一下采样模块相连，以复用所述第一下采用模块输出的特征，所述方法包括：将训练数据输入所述图像处理模型，得到所述图像处理模型的训练损失；根据所述训练损失，更新所述图像处理模型的模型参数；在更新所述模型参数的过程中，对所述第一下采样模块和所述第一上采样模块之间的网络架构进行神经网络架构搜索，以确定所述第一下采样模块和所述第一上采样模块之间的注意力模块的结构。

作为一种可能的实现方式，所述第一下采用模块和所述第一上采样模块通过多个注意力模块相连，所述多个注意力模块具有各自对应的连接权重，所述神经网络架构搜索用于更新所述多个注意力模块各自对应的连接权重。

作为一种可能的实现方式，所述方法还包括：根据所述多个注意力模块各自对应的连接权重的变化趋势，从所述多个注意力模块中选择目标注意力模块，作为所述U型网络训练完成之后需要保留的注意力模块，其中所述目标注意力模块对应的连接权重在所述网络结构搜索过程中呈上升趋势。

作为一种可能的实现方式，所述多个注意力模块包括以下注意力模块中的部分或全部：通道注意力模块、空间注意力模块和混合注意力模块。

作为一种可能的实现方式，所述图像处理模型为图像降噪模型，所述图像降噪模型包括噪声估计子网络和非盲去噪自网络，所述U型网络为所述非盲去噪子网络。

第二方面，一种图像处理模型的训练装置，所述图像处理模型包括U型网络，所述U型网络包括下采样通路和上采样通路，所述下采样通路包括第一下采样模块，所述上采样通道包括第一上采样模块，所述第一上采样模块与所述第一下采样模块相连，以复用所述第一下采用模块输出的特征，所述装置包括：输入模块，用于将训练数据输入所述图像处理模型，得到所述图像处理模型的训练损失；更新模块，用于根据所述训练损失，更新所述图像处理模型的模型参数；搜索模块，用于在更新所述模型参数的过程中，对所述第一下采样模块和所述第一上采样模块之间的网络架构进行神经网络架构搜索，以确定所述第一下采样模块和所述第一上采样模块之间的注意力模块的结构。

作为一种可能的实现方式，所述第一下采用模块和所述第一上采样模块通过多个注意力模块相连，所述多个注意力模块具有各自对应的连接权重，所述神经网络架构搜索用于更新所述多个注意力模块各自对应的连接权重。

作为一种可能的实现方式，所述装置还包括：选择模块，用于根据所述多个注意力模块各自对应的连接权重的变化趋势，从所述多个注意力模块中选择目标注意力模块，作为所述U型网络训练完成之后需要保留的注意力模块，其中所述目标注意力模块对应的连接权重在所述网络结构搜索过程中呈上升趋势。