[发明专利]一种U-net模型的轻量化方法及装置

说明书

本申请涉及深度学习技术领域，尤其涉及一种U‑net模型的轻量化方法及装置。该轻量化方法包括：获取U‑net模型；在所述U‑net模型中增加预裁剪因子层，得到第一模型；利用第一样本数据训练所述第一模型得到第二模型；对所述第二模型中预裁剪因子层的参数小于阈值的通道进行剪枝，得到第三模型。本申请实施例通过对U‑net模型进行剪枝，在加速网络模型推理的同时，保证了网络模型的准确率。

技术领域

本申请涉及深度学习技术领域，尤其涉及一种U-net模型的轻量化方法及装置。

背景技术

随着深度学习的发展，神经网络模型被应用到越来越多的应用场景中。但是，大型的神经网络模型的计算量大，需要耗费很多的计算资源，在计算量低的边缘移动设备中，耗时严重。

以人像抠图应用场景为例进行说明，在该应用场景中，经常使用神经网络模型作为图像特征的提取模块，神经网络模型对图像的表征更好从而可以取得良好的抠图效果，但是神经网络模型计算量大，需要耗费更多的计算资源，在计算量低的边缘移动设备中，耗时严重，难以实现实时人像抠图。

发明内容

有鉴于此，本申请实施例提供了一种U-net模型的轻量化方法及装置，可以解决相关技术中的至少一个技术问题。

第一方面，本申请一实施例提供了一种U-net模型的轻量化方法，包括：获取U-net模型；在所述U-net模型中增加预裁剪因子层，得到第一模型；利用第一样本数据训练所述第一模型得到第二模型；对所述第二模型中预裁剪因子层的参数小于阈值的通道进行剪枝，得到第三模型。

本实施例提供的一种U-net模型的轻量化方法，对U-net模型增加预裁剪因子层，并基于经训练的预裁剪因子层的参数进行剪枝，实现了网络模型的轻量化，在加速网络模型推理的同时，可以保证较高的准确率。

在一些实施例中，所述在所述U-net模型中增加预裁剪因子层，得到第一模型，包括：在所述U-net模型的编码器中各卷积模块的第一个卷积单元后面增加预裁剪因子层，得到第一模型。

在一些实施例中，在所述利用第一样本数据训练所述第一模型得到第二模型的反向传播过程中，利用损失函数的梯度更新所述第一模型中属于所述U-net模型的各层的参数，利用所述损失函数的梯度和惩罚项的梯度更新所述第一模型中预裁剪因子层的参数。

在一些实施例中，所述对所述第二模型中预裁剪因子层的参数小于阈值的通道进行剪枝，得到第三模型，包括：步骤S1，对所述第二模型中预裁剪因子层的参数小于阈值的通道进行剪枝，得到第三模型；步骤S2，若所述第三模型不满足预设条件，则更新所述阈值，并返回执行上述步骤S1；步骤S3，若所述第三模型满足预设条件，则存储所述第三模型。

在一些实施例中，所述对所述第二模型中预裁剪因子层的参数小于阈值的通道进行剪枝，得到第三模型之后，还包括：将所述第三模型中保留下来的预裁剪因子层与其前面的卷积单元进行融合，得到第四模型。

在一些实施例中，还包括：利用第二样本数据对所述第三模型进行调优训练，得到第四模型，所述第二样本数据的数量小于所述第一样本数据的数量，所述第二样本数据为与应用场景关联的样本数据。

在一些实施例中，所述第一样本数据为带标注的人像抠图数据，所述第三模型为用于人像抠图的模型。

第二方面，本申请一实施例提供了一种U-net模型的轻量化装置，包括：获取模块，用于获取U-net模型；增加模块，用于在所述U-net模型中增加预裁剪因子层，得到第一模型；训练模块，用于利用第一样本数据训练所述第一模型得到第二模型；剪枝模块，用于对所述第二模型中预裁剪因子层的参数小于阈值的通道进行剪枝，得到第三模型。

第三方面，本申请一实施例提供一种电子设备，包括：存储器、处理器以及存储在所述存储器中并可在所述处理器上运行的计算机程序，所述处理器执行所述计算机程序时实现如第一方面任一实施例所述的U-net模型的轻量化方法。