基于知识蒸馏的单幅图像去雾方法

基于知识蒸馏的单幅图像去雾方法

期刊：计算机工程 C

时间：2022

研究院：南京邮电大学

关键词：图像去雾；生成对抗网络；知识蒸馏；教师网络；学生网络

相关工作

基于学习的方法通过利用卷积神经网络（CNN） 从大量训练数据中提取图像特征，学习图像对之间的映射关系。早期的大多数方法仍然基于大气退化模型恢复无雾图像

例如，Ren 等[4]提出的 MSCNN 和 Cai 等[5]提出的 DehazeNet 是最早采用 CNN 进行图像去雾的算法，这两种方法都先使用 CNN 来估计透射率，再根据大气散射模型得到无雾图像。但由于缺乏有效的先验知识，难以准确估计大气光值 A 和传输率 t(x)，去雾效果并不理想

因此，第二类端到端的学习方法逐渐被提出，将有雾图像输入网络模型直接得到无雾图像，无需估计大气散射模型的中间参数

例如，Qu 等[6]将去雾任务视为图像到图像的转换问题，提出了一种基于增强的 pix2pix 去雾网络 EPDN，获得了具有良好视觉效果的无雾图像。Liu 等 [7]最近提出 GMAN，一种基于残差模块和编码器解码器结构的去雾模型,该模型在一些标准的去雾数据集上取得了良好的去雾效果

本文

1）设计一种基于编码器-解码器结构的条件生成对抗网络（CGANs）用于图像去雾。将蒸馏模型中的教师网络设计成具有较多网络层数及权值参数的重型网络，并设计三种不同参数规模且精简于教师网络的学生网络模型，以对比不同简化程度下学生网络的去雾性能。

（2）提出一种融合图像高、低频分量的方法（HLF -CGAN）， 即将频率信息作为去雾的附加约束条件。 首先将欲输入判别器的图片进行傅里叶变换，再使用拉普拉斯算子、高斯滤波器得到相应的滤除高频或低频后的图片，最后将三者拼接后一起输入判别器。

(3) 提出一种将以过程为导向的损失函数和以结果为导向的损失函数联合成总损失函数的方法来衡量教师网络和学生网络之间的特征相似性。

(4) 在室外模糊图像公开数据集 OTS 和室内模糊图像公开数据集 ITS 上进行网络的训练和测试。对比不同参数规模的学生网络去雾性能的变化，同时对比训练过程中相同迭代次数下教师网络和学生网络的去雾效果，并将本文提出的方法与其他去雾算法进行对比，目的是证明知识蒸馏在模型训练速率和图像去雾中的有效性。

本文提出的基于条件生成对抗网络和知识蒸馏的单幅图像去雾算法，网络的整体架构包括三个部分：教师网络（Teacher）、学生网络(Student)和蒸馏结构

教师网络

Mirza 等[13]提出了条件生成对抗网络（CGAN）， 它在生成器和判别器的输入端增加了额外的限制条件，使用源域图像代替噪声信息指导生成器生成目标图像，这在很多领域取得了优异的成果。本文将原始雾图作为条件信息（Condition）添加到判别器的输入端，从而判别器既要判断输入图像是否真实，同时也要比较输入图像与作为条件信息的原始雾图是否匹配。设计的基于 HLF-CGAN 的去雾算法中教师网络结构如图 1 所示。

图中 Generator 基于 U-Net 网络[8]进行改进， U-Net 网络是一个由编码器-解码器结构组成的 U 型对称网络，最初用于图像分割领域。编码器由 4 个编码块组成，每个编码块包含 2 个 3×3 大小的卷积核和 1 个 2×2 大小的最大池化层用于下采样，解码器包含 4 次反卷积和 4 次上采样过程。虽然采用最大池化的方式进行下采样不需要额外增加参数，但容易丢失高频特征，因此本文采用卷积步长为 2 的下采样方式，以减少特征的丢失。从图 1 的网络结构可以看出生成器网络是由编码器、解码器和跳跃连接层组成。编码器部分经过 8 次卷积操作将大小为 256 ×256×3 的输入图像最终转换成大小为 1×1× 512 的特征图，编码器的主要作用是提取图像特征信息。具体地，每一层的卷积步长为 2，卷积核大小为 5×5，卷积核个数前三层分别为 64、128、256 个，后五层均为 512 个。

解码器部分在此基础上进行 8 次反卷积操作， 对应位置采用同样大小和数量的卷积核，直到最终输出大小为 256×256×3 的特征图。跳跃连接结构将编码器和解码器中相同大小的特征图按通道拼在一起，以更好地利用在下采样过程中保留的局部细节信息，加快网络训练过程。

高频和低频分量是图像的基本特征，去雾图像和对应的真实无雾图像的高、低频分量也会存在一定差异。高频分量包括图像中强度变化迅速的部分，如锐利的边缘、纹理和精细的细节。相反，低频分量一般是图像强度值变化缓慢的部分，更加强调图像的亮度、颜色和对比度信息[14]。将高频和低频信息同时作为判别网络的约束条件，有利于判别器对图像的边缘、纹理和颜色上差异性的判断。对比实验如图 2 所示，图片下方记录了去雾量化指标。从图中可以看出加入频率信息作为附加条件的网络（HLF-CGAN） 比直接将原图输入判别器的网络（CGAN）表现出更好的去雾效果，去雾评价指标取得更高的值。 图 1 中 Discriminator 采用 PatchGAN 结构， PatchGAN 与普通 GAN 的主要区别在于普通 GAN 的判别器是将输入数据直接映射成一个概率值，而 PatchGAN 是将输入数据映射成一个 N×N 大小的矩阵，矩阵中每一个元素代表着输入图像中一个较大的感受野。该矩阵的每一个元素只有“0”或“1” 两个值，即该元素为真的概率，最后再求取平均值作为判别器的最终输出。PatchGAN 的优点是可以从原输入图像较大的感受野中关注到更多图像的细节信息。判别器对输入图像经过 4 次卷积操作转换成大小为 16×16×512 的特征图，再将该大小的特征图拉伸成一个一维矩阵，通过 sigmoid 函数得到每一个元素的概率值后求取平均值作为最终判别器的输出。 本文对生成器的输出使用 L1 正则项（最小绝对值偏差）来约束生成图像与真实图像之间的距离，使生成图像的特征分布更加逼近真实图像的特征分布。 与 L2 正则项（最小均方误差）相比，L1 正则项具有更好的鲁棒性。因此，教师网络的总损失函数可以被定义为

学生网络

本文设计了三个与教师网络结构相似且具有不同参数规模的学生网络，分别记为 Student-1、 Student-2、Student-3。经计算，教师网络共包含 3052.30 万个权值参数，模型大小为 509852kb。 Student-1 网络的生成器由 6 层下采样和 6 层上采样的编码器-解码器结构组成，包含 1741.47 万个权值参数，模型大小为 197367kb，约为教师网络的 1/2 倍；其余两个学生网络的生成器分别由 5 层、4 层上下采样的编码器-解码器结构组成，模型大小分别约为教师网络的 1/5、1/10 倍，具体参数如表 1 所示。从表中可以看出，相较于教师网络的 8 层卷积与反卷积结构，学生网络具有更少的权值参数，占用更少的内存空间，这显著的节省了大量的计算资源。

知识蒸馏

Romero 等人[15]提出的知识蒸馏方法仅让学生网络简单的模仿教师网络的中间结果，考虑到教师网络和学生网络两层特征图之间的关系，本文首先通过最大池化将两层卷积层转化成相同的尺寸，再将其对应的通道两两作内积得到 FSP 矩阵，即图 3 中的 T G和 S G 。FSP 矩阵是由 Yim[16]等提出的一种用于解决问题的流，若将流定义为两个不同层的特征图的关系， 知识蒸馏则为解决问题的流的方法。从图中可以看出，本文教师网络的两个 FSP 矩阵是由教师网络的第 1 层和第 3 层、第 4 层和 6 层特征图分别两两作内积得到。同理，学生网络的两个 FSP 矩阵也是由网络的相同层特征图作内积得到。得到教师网络和学生网络的 FSP 矩阵后，再利用 L2 损失函数去拉近两个网络 FSP 矩阵之间的距离，以达到知识蒸馏的目的。FSP 矩阵的定义如下