[发明专利]一种深度学习图像编解码网络的蒸馏训练方法与系统

说明书

本发明公开了一种深度学习图像编解码网络的蒸馏训练方法与系统，对深度学习图像编解码网络做轻量化蒸馏操作，保证了在压缩性能基本不下降的情况下，有效降低网络参数量70％，在实际生产应用中，经过蒸馏学习的学生网络占用更少的设备内存，拥有更快的编解码推理速度，可以在中心端，边缘端设备上编解码高清图像。

技术领域

本发明涉及图像编解码技术领域，尤其涉及一种深度学习图像编解码网络的蒸馏训练方法与系统。

背景技术

目前学术界基于深度学习的图像编解码网络，在基于PSNR(峰值信噪比)，MS-SSIM(多层级结构相似性)和人眼主观的压缩性能上均超越了传统H266，H265，WEBP，JPEG的效果。

但是，现有基于深度学习的图像编解码网络的编码和解码计算量巨大，难以真正的应用于在实际生产中。

发明内容

本发明的目的是提供一种深度学习图像编解码网络的蒸馏训练方法与系统，在保证压缩性能基本不下降的情况下，极大的减小了编解码过程的内存占用和计算量。

本发明的目的是通过以下技术方案实现的：

一种深度学习图像编解码网络的蒸馏训练方法，包括：

第一阶段：构建包含教师编解码模型的图像编解码网络，所述教师编解码模型包括：教师编码模型、超先验概率模型与教师解码模型，所述教师编码模型与教师解码模型构成教师编解码模型；输入原始图像至所述包含教师编解码模型的图像编解码网络，计算第一阶段的率失真损失，并对所述基于教师编解码模型的图像编解码网络进行训练；

第二阶段：在所述包含教师编解码模型的图像编解码网络中添加学生编解码模型，所述学生编解码模型包括学生编码模型与学生解码模型，所述学生编解码模型的位置与所述教师编解码模型的位置相同；所述学生编解码模型通过将教师编解码模型中的残差模块做轻量化处理后得到；输入原始图像至所述教师编码模型与学生编码模型，教师编码模型与学生编码模型的输出均经所述超先验概率模型后，对应的进入教师解码模型与学生解码模型；利用教师编码模型与学生编码模型的输出、教师编码模型与学生编码模型中间相同位置的输出、以及教师解码模型与学生解码模型中间相同位置的输出计算知识蒸馏损失，并结合学生解码模型的输出与超先验概率模型的输出计算的第二阶段率失真损失对学生编解码模型进行训练，更新学生编解码模型的参数；训练达到指定次数后，移除教师编解码模型，输入原始图像，计算第二阶段率失真损失，并对所述学生编解码模型与超先验概率模型的参数进行微调。

一种深度学习图像编解码网络的蒸馏训练系统，包括：

第一训练单元，应用于第一阶段，所述第一阶段包括：构建包含教师编解码模型的图像编解码网络，所述教师编解码模型包括：教师编码模型、超先验概率模型与教师解码模型，所述教师编码模型与教师解码模型构成教师编解码模型；输入原始图像至所述包含教师编解码模型的图像编解码网络，计算第一阶段的率失真损失，并对所述基于教师编解码模型的图像编解码网络进行训练；

第二训练单元，应用于第二阶段，所述第二阶段包括：在所述基于教师编解码模型的图像编解码网络中添加学生编解码模型，包括学生编码模型与学生解码模型，所述学生编解码模型的位置与所述教师编解码模型的位置相同；所述学生编解码模型通过将教师编解码模型中的残差模块做轻量化处理后得到；输入原始图像至所述教师编码模型与学生编码模型，教师编码模型与学生编码模型的输出均经所述超先验概率模型后，对应的进入教师解码模型与学生解码模型；利用教师编码模型与学生编码模型的输出、教师编码模型与学生编码模型中间相同位置的输出、以及教师解码模型与学生解码模型中间相同位置的输出计算知识蒸馏损失，并结合学生解码模型的输出与超先验概率模型的输出计算的第二阶段率失真损失对学生编解码模型进行训练，更新学生编解码模型的参数；训练达到指定次数后，移除教师编解码模型，输入原始图像，计算第二阶段率失真损失，并对所述学生编解码模型与超先验概率模型的参数进行微调。