一种基于特征知识蒸馏的轻量级图像去噪模型

摘""要："为了构建适用于小型设备的轻量级图像去噪模型，提出了一种基于特征知识蒸馏的新方法. 该方法通过学习教师模型的特征图，捕捉深层知识，从而构建轻量级去噪模型，其参数量仅为教师模型的五分之一. 实验结果验证了蒸馏算法的有效性，在不同噪声水平及数据集下，都显著提升了学生模型的去噪性能，为轻量级图像去噪模型的构建提供了一种新的方向.

关键词："知识蒸馏；"特征学习；"卷积神经网络；"图像去噪

中图分类号："TP 183 """文献标志码："A """文章编号："1000-5137（2024）04-0489-07

A lightweight image denoising model based on feature knowledge distillation

SHEN Yu，"ZHANG Dingni

（College of Information，"Mechanical and Electrical Engineering，"Shanghai Normal University，"Shanghai 201418，"China）

Abstract："To construct a lightweight image denoising model suitable for small-scale devices，"a novel approach based on feature knowledge distillation was proposed in this paper. Deep-seated knowledge within a teacher model was captured through learning from its feature maps by this method，"resulting in the creation of a lightweight denoising model with parameters only one-fifth the size of the teacher model. Experimental results validated the effectiveness of the distillation algorithm，"demonstrating significant improvements in denoising performance for the student model across varying noise levels and datasets，"which introduced a promising avenue for constructing lightweight image denoising models.

Key words："knowledge distillation；"feature learning；"convolution neural network；"image denoising

图像去噪作为一种经典的低级视觉任务，在医学影像、卫星图像、地质图像等领域均有重要作用. 一般来说，图像去噪方法主要分为两类：以BM3D^［1］为代表基于模型的方法和以DnCNN^［2］为代表基于判别学习的方法. 前者通常需要手工制作的图像先验，例如稀疏性^［3］和非局部自相似性^［4］，这些图像先验可能不够强，无法表征复杂的图像结构；而后者凭借残差学习、批归一化（BN）以及卷积神经网络（CNN）强大的建模能力，具有优秀的去噪性能，并衍生出了一系列性能优异的图像去噪网络^［5-7］. 但深度网络具有庞大的参数量，对设备的计算要求较高，所以如何构建轻量级的图像去噪模型成为一个重要的研究方向.

知识蒸馏最初由HINTON等^［8］提出，其主要思想是通过一个小型模型（参数量较小、性能较弱的模型，称为学生模型）学习一个大模型（参数量较大、性能较强的模型，称为教师模型）的知识，使得学生模型能凭借很小的参数量，取得接近于教师模型的性能. 知识蒸馏的思想在分类任务的模型压缩中运用广泛，主要可以分为logits向量蒸馏和特征蒸馏^［9］.由于logits向量蒸馏无法学习到教师模型内部的丰富特征，其性能一般不如特征蒸馏. 另外，logits向量蒸馏需要学习logits向量层的输出分布，这对于图像去噪任务来说，与直接学习输出损失无异.

本文作者借助特征知识蒸馏的思想，将高性能的教师模型特征迁移到学生模型中，提出了一个轻量的图像去噪模型，分析了特征蒸馏在较高噪声条件下迁移特征的弊端，证明了知识蒸馏对于压缩去噪模型的有效性.

1 "模型介绍

1.1 DnCNN模型

DnCNN模型的主要思想是将图像中的噪声提取出来，利用残差学习和BN处理来训练更深层次的网络，其网络架构如图1所示.

1.2 特征知识蒸馏算法

1.3 蒸馏模型介绍

学生模型和教师模型使用相同的网络架构，便于更好地拟合特征损失. 模型蒸馏的工作流程如图2所示.

虽然以Conv+BN+ReLU组成的模块确实可以表达更丰富的特征，但是过深的网络中间层会产生一些多余的知识，不仅会占用模型的内存，还会降低模型运行时的速度，因此需要对模型进行压缩. 利用特征知识蒸馏来构建更为轻量的学生模型，使其学习教师模型庞大的特征，进而将教师模型丰富的特征迁移到学生模型中，使学生模型可以获得与教师模型相接近的性能.

2 "实验结果及分析

2.1 实验参数

实验使用DIV2K^［11］数据集，总共800张图片，按照7∶1的比例划分训练集和验证集，将每张图片随机裁剪成4块40×40的patch，每个patch顺时针旋转90°，180°和270°用于数据增强，训练集数量总共为11 200. 测试集使用BSD68^［12］和Set12^［2］，测试集的输入尺寸为256×256. 为了在已知噪声水平下，训练模型的去噪性能，考虑了三种加性高斯白噪声，即"= 15，25和50. 教师模型和学生模型的架构均为DnCNN^［2］，教师网络的层数为48，学生网络为12. 硬损失和特征损失均采用均方差损失函数（MSE），优化器为Adam，学习率固定1×10^-4，batch_size为128，训练500轮. 蒸馏训练与模型训练的参数设置相同. 训练设备：GPU是RTX 4 090（24 GB），CPU是"Intel（R）"Xeon（R）"Platinum 8352 V CPU @ 2.10 GHz；训练环境：PyTorch 2.0.0，Python 3.8（ubuntu20.04），Cuda 11.8.

2.2 评价指标

选取最常用的两个评估指标来衡量模型的去噪性能：峰值信噪比（PSNR）和结构相似性指数（SSIM）^［13］. 其中，PSNR可以较为准确地衡量图像之间的失真程度，PSNR值越大，表示噪声图像和干净图像之间的失真越小，即图像的质量越好. SSIM旨在测量两幅图像之间的结构相似性，其值为0～1，越接近1，表明两个图像在结构上越相似，即噪声图像越接近于干净图像.

2.3 实验结果

特征蒸馏的训练过程如图3所示. 可以发现，训练过程中，验证集和训练集的拟合程度很高，两者的PSNR曲线几乎重合，这表明模型训练没有发生过拟合或者欠拟合的情况.

实验结果如表1和2所示. 从表1的PSNR指标可以看出，不同数据集及噪声水平下，知识蒸馏都可以在一定程度上提升模型的去噪性能. 然而在较高噪声水平下，受限于参数量太少，模型对噪声的拟合程度有限，去噪能力依旧有待提升. 同样地，表2的SSIM指标也验证了特征蒸馏对模型去噪能力的提升效果较为明显.

从图4可以看出，学生模型在修复图片后依然会残留不少噪声，相较之下，蒸馏模型的残留噪声就少了很多，非常接近教师网络的修复效果，尤其是在边缘部分以及色层连续的部分，如：天空、房屋表面、水果表皮等.

2.4 消融实验

分别计算了学生模型在特征蒸馏前后的损失值和验证集上的PSNR指标，结果如图5所示.

图5表明，学生模型在蒸馏后，损失下降更为明显，并且模型收敛后的损失更小，这一点在验证集的PSNR曲线上也表现得十分明显，说明模型蒸馏后，去噪性能更好，收敛后更加稳定，性能波动更小. 通过知识蒸馏来学习教师模型的特征，可以显著提升原模型的性能.

2.5 性能比较

在实际应用中，对于轻量级图像去噪模型的要求不仅在于去噪效果良好，还要具备更快的推理速度，使其能应用于小型移动设备. 表5为教师模型和蒸馏模型的参数量以及推理时间对比. 对于模型推理时间的计算只在CPU上进行. 可以发现，蒸馏模型的参数量不到教师模型的1/5，推理时间未超过教师模型的1/4.

3 "结论

采用知识蒸馏方法，将教师模型的隐含特征知识传递给学生模型，在显著减小参数量的同时，使学生模型能获得接近于教师模型的去噪性能.为验证这一理论，进行了广泛的消融实验，结果表明：特征知识蒸馏在不同噪声水平和数据集下均能提升模型的PSNR和SSIM指标. 在性能比较方面，所提出的轻量级去噪模型超越了近两年来最新的模型，充分证明了知识蒸馏在去噪模型压缩领域的有效性.

参考文献：

［1］ DABOV K，"FOI A，"KATKOVNIK V，"et al. Image denoising by sparse 3-D transform-domain collaborative filtering ［J］. IEEE Transactions on Image Processing，"2007，16（8）：2080-2095.

［2］ ZCOZHANG K，"ZUO W M，"CHEN Y J，"et al.，"Beyond a Gaussian denoiser："residual learning of deep CNN for image denoising ［J］. IEEE Transactions on Image Processing，"2017，26（7）：3142-3155.

［3］ ELAD M，"AHARON M. Image denoising via sparse and redundant representations over learned dictionaries ［J］. IEEE Transactions on Image Processing，"2006，15（12）：3736-3745.

［4］ BUADES A，"COLL B，"MOREL J M. A non-local algorithm for image denoising ［C］// IEEE Computer Society Conference on Computer Vision and Pattern Recognition. San Diego："IEEE，"2005：60-65.

［5］ LIANG Y Z，"LIANG W. ResWCAE："biometric pattern image denoising using residual wavelet-conditioned autoencoder ［J/OL］. ［2023-05-08］. arXiv：2307.12255，"2023. https：//arxiv.org/abs/2307.12255.

［6］ HUANG Z H，"ZHU Z F，"WANG Z C，"et al. DGDNet："deep gradient descent network for remotely sensed image denoising ［J］. IEEE Geoscience and Remote Sensing Letters，"2023，20：1-5.

［7］ ZHANG Q，"XIAO J Y，"TIAN C W，"et al. A robust deformed convolutional neural network （CNN）"for image denoising ［J］. CAAI Transactions on Intelligence Technology，"2023，8（2）：331-342.

［8］ HINTON G，"VINYALS O，"DEAN J. Distilling the knowledge in a neural network ［J/OL］. ［2023-05-08］.arXiv：1503.02531，"2015. https：//www.arxiv.org/abs/1503.02531.

［9］ ZAGORUYKO S，"KOMODAKIS N. Paying more attention to attention："improving the performance of convolutional neural networks via attention transfer ［J/OL］. ［2023-05-08］. arXiv：1612.03928，"2016. https：//arxiv.org/abs/1612. 03928v2.

［10］ IOFFE S，"SZEGEDY C. Batch normalization："accelerating deep network training by reducing internal covariate shift ［C］// International Conference on Machine Learning. Lille："JMLR，"2015：448-456.

［11］ TIMOFTE R，"AGUSTSSON E，"VAN GOOL L，"et al. Challenge on single image super-resolution："methods and results ［C］// Conference on Computer Vision and Pattern Recognition. Salt Lake City："IEEE，"2018：852-863.

［12］ MARTIN D，"FOWLKES C，"TAL D，"et al. A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics ［C］// Proceedings Eighth IEEE International Conference on Computer Vision. Vancouver："IEEE，"2001：416-423.

［13］ WANG Z，"BOVIK A C，"SHEIKH H R，"et al. Image quality assessment："from error visibility to structural similarity ［J］. IEEE Transactions on Image Processing，"2004，13（4）：600-612.

［14］ GUO M H，"XU T X，"LIU J J，"et al. Attention mechanisms in computer vision："a survey ［J］. Computational visual media，"2022，8（3）：331-368.

［15］ JANJUŠEVIĆ N，"KHALILIAN-GOURTANI A，"WANG Y. CDLNet："robust and interpretable denoising through deep convolutional dictionary learning ［J/OL］. ［2023-05-08］. arXiv：2103.04779，"2021. https：// arxiv.org/abs/2103.04779.

［16］ JO Y，"CHUN S Y，"CHOI J. Rethinking deep image prior for denoising ［C］// International Conference on Computer Vision. Montreal："IEEE，"2021：5087-5096.

（责任编辑：包震宇）