何 涛,俞舒曼,徐 鹤
(1.南京邮电大学 电子与光学工程学院、微电子学院,南京 210023;2.南京邮电大学 计算机学院、软件学院、网络空间安全学院,南京 210023;3.江苏省无线传感网高技术研究重点实验室,南京 210023)
雾是由细尘或烟雾颗粒在大气中聚集得到的产物,在雾霾天气下采集的图片通常会失去对比度、颜色保真度和边缘信息,从而降低场景的可视性。例如,若环境中存在大量雾霾,在执行视频监控任务时将难以从视频中定位到有效的目标信息,进而阻碍任务的顺利执行。因此,去除图像中的雾霾并提高图像清晰度成为计算机视觉领域的研究热点。
近年来,研究人员提出基于先验知识和基于学习的去雾方法。基于先验知识的去雾方法是根据手工设计的特征,并结合文献[1]提出的大气散射模型以恢复清晰的无雾图像。文献[2]提出一种经典的统计观测方法,即在大多数非天空区域中,某些像素中至少有一个颜色通道的强度较低,从而根据暗通道先验知识可以粗略估计出透射率和大气光值。文献[3]是针对模糊图像的场景深度构建一种线性模型,根据颜色衰减来恢复深度信息,利用模糊图像的深度图估计透射率。因透射率估计不准确,导致在输出去雾图像中存在明显的伪影问题。
基于学习的去雾方法是利用卷积神经网络(Convolutional Neural Network,CNN)从大量训练数据中提取图像特征,同时学习图像对之间的映射关系。早期的大多数方法是基于大气退化模型恢复无雾图像,例如MSCNN[4]和DehazeNet[5]方法是最早采用CNN 进行图像去雾,再基于大气散射模型得到无雾图像。由于这两种方法缺乏有效的先验知识,因此难以准确估计大气光值和传输率,去雾效果并不理想。端到端的学习方法是将有雾图像输入到网络模型中直接得到无雾图像,无需估计大气散射模型的中间参数。例如,文献[6]将去雾任务看成是图像到图像的转换问题,构建一种基于增强的pix2pix去雾网络EPDN,获得具有较优视觉效果的无雾图像。文献[7]提出一种基于残差模块和编码器-解码器结构的去雾模型GMAN,该模型在一些标准的去雾数据集上具有较优的去雾效果。
本文提出一种基于条件生成对抗网络(Conditional Generative Adversarial Network,CGAN)与知识蒸馏的单幅图像去雾方法KD-GAN。该方法以改进的U-Net网络[8]作为生成器的基本架构,并在判别器的输入端融合图像高、低频分量,将频率信息作为去雾的附加约束条件与原始图像一起输入到判别器中。在此基础上,通过知识蒸馏的方法减小条件生成对抗网络的计算量和存储空间,从而加快模型的训练速度。
近年来,生成对抗网络[9](Generative Adversarial Network,GAN)在图像生成、图像插值和目标分类等领域中取得较大进展。GAN 由生成器和判别器两个部分组成。生成器G:Z→Y,将输入的随机噪声z∈Z映射到输出图像y∈Y中;判别器D:y→[0,1],将其输入的图像y映射成概率值。针对去雾网络模型,生成器的作用是利用给定的一组有雾图像{y1,y2,···,yn}生成伪无雾图像判别器是通过对比真实图像x与生成器生成的图像y*的特征分布,从而预测输入的图像为真实图像x的概率。生成器的优化目标是在保留输入图像的结构和细节信息的同时,尽可能地生成接近真实无雾图像特征分布的去雾图像,从而使得判别器无法辨别图像来源。判别器的优化目标则是尽可能地辨别出输入图像是生成的去雾图像还是真实无雾图像。生成器和判别器相互博弈,以达到纳什均衡[10]。GAN 的损失函数如式(1)所示:
其中:Pdata(x)为真实图像数据的分布;Pz(z)为输入生成器噪声z的分布;D(x)为判别器判断真实图像为真的概率;D(G(z))为判别器判断生成图像为真的概率。
由于生成器需要大量的参数来生成具有高维语义信息的图像,因此大多数GAN 的执行会消耗较多的计算资源,从而限制了GAN 在CPU 性能有限的移动设备以及边缘设备上的应用。近年来,研究人员对重型深度模型的压缩和加速方法进行研究。如文献[11]提出一种压缩技术,通过将教师网络产生的分类概率分布作为软目标来训练小模型,从而将集成模型中的知识迁移到单个模型中,这种压缩技术又称知识蒸馏。在知识蒸馏模型中将具有复杂网络结构、学习能力更强的重型网络视为教师网络,将网络结构简单、学习能力较弱的轻量型网络视为学生网络。文献[12]指出,浅层网络可以学习到深层网络的复杂函数,并实现以前只有深度模型才能达到的精度。受上述方法的启发,本文提出一种基于条件生成对抗网络和知识蒸馏的单幅图像去雾方法(KD-GAN),本文将蒸馏模型中的教师网络设计成具有较多网络层数及权值参数的重型网络,并设计3 种不同参数规模且精简于教师网络的学生网络模型,以对比在不同简化程度下学生网络的去雾性能。在室外模糊图像公开数据集OTS 和室内模糊图像公开数据集ITS 上进行网络的训练和测试。本文通过对比不同参数规模学生网络的去雾性能变化,同时对比训练过程中相同迭代次数下教师网络和学生网络的去雾效果,并将本文方法与其他去雾方法进行对比,验证知识蒸馏在模型训练速率和图像去雾中的有效性。
本文提出的基于CGAN 和知识蒸馏的单幅图像去雾方法,网络的整体架构包括教师网络、学生网络和知识蒸馏结构3 个部分。知识蒸馏[11]将提炼的教师网络的知识传递给结构简单、学习能力较弱的学生网络,在提高学生模型训练速度的同时,尽可能使其性能接近于具有复杂结构、学习能力强的教师网络。
文献[13]提出CGAN,其在生成器和判别器的输入端增加了额外的限制条件,通过源域图像代替噪声信息,以指导生成的目标图像。本文将原始有雾图像作为条件信息添加到判别器的输入端,判别器既要判断输入图像是否真实,同时判断输入图像与作为条件信息的原始有雾图像是否匹配。基于引入频率信息作为附加条件网络(HLF-CGAN)的去雾方法中教师网络结构如图1 所示。
图1 教师网络结构Fig.1 Structure of teacher network
图中生成器是基于U-Net 网络[8]进行改进,U-Net 网络是由编码器-解码器结构组成的U 型对称网络,最初用于图像分割领域。编码器由4 个编码块组成,每个编码块包含2 个3×3 大小的卷积核和1 个2×2 大小的最大池化层,用于下采样,解码器包含4 次反卷积和4 次上采样过程。虽然U-Net 网络采用最大池化的方式进行下采样,不需要额外增加参数,但容易丢失高频特征。因此本文采用卷积步长为2 的下采样方式,以减少特征的丢失。从图1 可以看出,生成器网络是由编码器、解码器和跳跃连接层组成。编码器部分经过8 次卷积操作将大小为256×256×3 的输入图像转换成大小为1×1×512 的特征图,其主要作用是提取图像特征信息。每层的卷积步长为2,卷积核大小为5×5,卷积核个数前3 层分别为64、128、256 个,后5 层均为512 个。
解码器部分进行8 次反卷积操作,在与编码器相对应的位置上采用同样大小和数量的卷积核,直到最终输出大小为256×256×3 的特征图。跳跃连接结构将编码器和解码器中相同大小的特征图按通道拼在一起,以更好地利用在下采样过程中保留的局部细节信息,从而加快网络训练过程。
高频和低频分量是图像的基本特征,去雾图像和对应的真实无雾图像的高、低频分量也会存在一定差异。高频分量包括图像中强度变化较快的部分,如锐利的边缘、纹理和精细的细节。相反,低频分量一般是图像强度值变化缓慢的部分,更加强调图像的亮度、颜色和对比度信息[14]。将高频和低频信息同时作为判别网络的约束条件,有利于判别器对图像边缘、纹理和颜色上的差异性进行判断。不同网络的去雾结果对比如图2 所示。
图2 不同网络的去雾结果对比Fig.2 Dehazing results comparison among different networks
图片下方为去雾量化指标(结构相似性SSIM/峰值信噪比PSNR)。从图2 可以看出,HLF-CGAN相比直接将原图输入到判别器的网络(CGAN)具有更优的去雾效果。
图1 中判别器采用PatchGAN 结构,PatchGAN与普通GAN 的主要区别在于:普通GAN 的判别器是将输入数据直接映射成一个概率值,而PatchGAN是将输入数据映射成一个N×N大小的矩阵,矩阵中每个元素代表输入图像中一个较大的感受野。该矩阵的每个元素只有0 或1,即该元素的概率为真,最后再将求取的平均值作为判别器的最终输出。PatchGAN 的优点是可以从原输入图像较大的感受野中提取图像更多的细节信息。判别器对输入图像进行4 次卷积操作,并转换成大小为16×16×512 的特征图,再将该特征图拉伸成个一维矩阵,通过sigmoid 函数得到每个元素的概率值后再求取平均值,并将其作为最终判别器的输出。
本文通过L1 正则项(最小绝对值偏差)约束生成图像与真实图像之间的距离,使得生成图像的特征分布更加接近真实图像的特征分布。与L2 正则项(最小均方误差)相比,L1 正则项具有更优的鲁棒性。教师网络的总损失函数如式(2)所示:
其中:Pdata(x)为输入生成器图像的数据分布;E(·)为期望值;x、y分别为原始有雾图像和真实无雾图像;D(x,y)、D(x,G(x))分别为判别器判断真实图像为真的概率及判断生成图像为真的概率;符号◦为拼接的含义;x◦xLF◦xHF为将原始雾图与其滤波后的低频、高频图像拼接;λ为可调节的超参数,代表损失项的权重。
虽然教师网络具有更多权值参数和网络层数,能够优化去雾效果,但是其模型规模庞大,导致预训练耗时久,从而难以应用到实际场景中。例如,在自动驾驶领域中,当图像识别系统遇到雾天路况时,为达到安全驾驶的目的,系统需要及时进行去雾处理,此时自动驾驶的安全性将与系统的响应速度、去雾性能有着密切的关系。因此,如何压缩模型规模并提高模型训练速度成为研究热点。
本文设计了3 个与教师网络结构相似,但具有不同参数规模的学生网络,分别记为学生网络-1、学生网络-2、学生网络-3。经计算,教师网络共包含3 052.30×104个权值参数,模型大小为509 852 KB。学生网络-1 的生成器由6 层下采样和6 层上采样的编码器-解码器结构组成,包含1 741.47×104个权值参数,模型大小为197 367 KB,约为教师网络的1/2。其他2 个学生网络的生成器分别由5 层和4 层的上下采样编码器-解码器结构组成,模型大小分别约为教师网络的1/5 和1/10,与教师网络的8 层卷积与反卷积结构相比,学生网络具有更少的权值参数,并占用更小的内存空间,从而节省大量的计算资源。
本文提出的知识蒸馏网络结构如图3 所示,其中学生网络结构以学生网络-1 为例。
图3 知识蒸馏结构Fig.3 Structure of knowledge distillation
文献[15]提出的知识蒸馏方法是学生网络简单模仿教师网络的中间结果,考虑到教师网络和学生网络两层特征图之间的关系,本文首先通过最大池化将两层卷积层转化成相同的尺寸,再将其对应的通道两两作内积得到FSP 矩阵,即图3 中的GT和GS。FSP 矩阵是由Yim 等[16]提出的一种用于解决问题的流,若将流定义为2 个不同层的特征图关系,知识蒸馏则是解决问题的流的方法。从图3 可以看出,本文教师网络的2 个FSP 矩阵是由教师网络第1 层、第3 层、第4 层和第6 层的特征图分别两两作内积得到。同理,学生网络的2 个FSP 矩阵也是由网络相同层的特征图作内积得到。本文在得到教师网络和学生网络的FSP 矩阵后,再利用L2 损失函数缩短这2 个网络FSP 矩阵之间的距离,以达到知识蒸馏的目的。FSP 矩阵的定义如式(3)所示:
其中:x、W分别为输入图像和网络的权值参数;为网络中第i层特征图;h、w、m分别为特征图的高度、宽度和通道数为网络中第j层特征图。考虑到L2 正则项对异常点具有放大效果,本文选择L2 范数的平方作为约束2 个网络FSP 矩阵的代价函数,因此中间层知识蒸馏的损失函数如式(4)所示:
其中:N为损失函数的个数;λi为损失项的权值,取值为10~100[17]。通常λ值越大,代价函数在参数为0 时取到最小值,因为正则化系数越大,正则化的函数图形会加速向坐标轴原点收缩。本文在基于λ为固定值100 时进行实验。本文提出联合损失函数来衡量教师网络和学生网络之间的特征相似性,即同时使用L1 正则项对2 个网络生成器的输出端特征图进行约束。因此,学生网络生成器的总损失函数如式(5)所示:
其中:GT(x)、GS(x)分别为教师网络和学生网络生成器输出的特征图。
本文实验主要分为以下5 个步骤:
1)首先冻结学生网络权值参数,开始训练教师模型,学生网络不参与训练,直到教师网络的生成器和判别器的损失函数收敛。损失函数如式(2)所示。
2)当迭代第1×104、2×104、3×104、5×104次时,记录教师网络训练过程中测试集图像的去雾结果。
3)冻结并保存预训练好的教师网络的权值参数,在相同的数据集上开始训练学生网络,此时教师网络不参与训练,直到学生网络的生成器和判别器的损失函数收敛。损失函数如式(5)所示。
4)记录学生网络在训练过程迭代相同次数时测试集图像的去雾结果。
5)对比教师网络与学生网络在迭代相同次数时测试集中图像的去雾效果,得到相应的实验结论。
本文从室外合成雾天图像数据集(OTS)和室内合成图像数据集(ITS)中分别选取5 000 张组合成模型的训练集。在OTS数据集上选取不重复的1 000张雾图和室外真实雾天图像(HSTS)数据集中选取1 000张雾图组合成模型的测试集。
本文方法在NVIDIA RTX 1080Ti GPU 计算机的Ubuntu16.0 系统上运行,并使用Tensorflow1.4.0深度学习框架进行实验。生成器输入图像大小设置为256×256×3,学习率设置为固定值0.000 2,动量参数设置为0.5,采用Adam 优化器对网络进行优化。
为客观地评估去雾效果,本文使用峰值信噪比PSNR[18]和结构相似性SSIM[19]作为去雾评价指标。PSNR 值取决于原始雾图与生成图像对应像素间的误差,其值越大,恢复得到的无雾图像清晰度越高[20]。SSIM 从亮度、对比度、结构3 个方面对图像相似性进行评估,其值越大,代表图像失真越小。
3.2.1 不同参数规模的学生网络去雾性能对比
图4 为教师网络与3 种不同参数规模学生网络的去雾结果对比,记为A 组和B 组。表1 为教师网络与不同参数规模学生网络的评价指标与网络规模对比。学生网络-1 在具有更少模型参数的情况下,可以达到与教师网络相近的去雾性能;学生网络-2 和学生网络-3 的去雾性能呈现较明显的下降趋势。从图4 中方框选出的部分可以看出,学生网络-3 对远景及细节部分恢复得比较模糊,失真较严重。实验结果表明,本文基于蒸馏架构,当学生网络参数规模简化为教师网络的1/2 时,能够达到最佳的蒸馏效果,当继续简化学生网络参数规模时,网络的去雾评价指标与视觉感受都会逐渐下降。
表1 教师网络与学生网络的评价指标与网络规模对比Table 1 Evaluation indexs and network scales comparison of teacher network and student network
图4 教师网络与学生网络的去雾结果对比Fig.4 Dehazing results comparison of teacher network and student network
3.2.2 知识蒸馏在模型收敛速率中的有效性
本文对不同迭代次数下教师网络和学生网络(学生网络-1)的测试结果进行分析,分别记录了迭代第1×104、2×104、3×104和5×104次时2 个网络的去雾结果。图5 是教师网络和学生网络在迭代第1×104次和第2×104次时生成器输出的去雾图像与原始雾图的对比。
图5 在不同迭代次数下教师网络和学生网络的去雾结果对比Fig.5 Dehazing results comparison of teacher networks and student networks under different iteration times
从迭代了相同次数的教师网络和学生网络的去雾结果可以看出,学生网络具有更优的去雾性能。在不同迭代次数下教师网络和学生网络的去雾评价指标如表2 所示。
表2 在不同迭代次数下教师网络和学生网络去雾评价指标对比Table 2 Dehazing evaluation indexes comparison of teacher networks and student networks under different iteration times
从表2 可以看出,在迭代第1×104和2×104次时,学生网络的PSNR、SSIM 值略优于迭代相同次数时的教师网络,并且学生网络在迭代到第3×104次时PSNR、SSIM 数值已接近于教师网络迭代第5×104次时的值。这是由于训练好的重型网络通常具有大量信息,轻量型网络通过重型网络的知识迁移优化权值参数时,只需要对参数进行微调,因此,轻量型网络的训练速度远优于从头开始训练的网络。
由此可以得出:迭代次数的增加对于网络性能的提升起着至关重要的作用,由于学生网络得到了从教师网络迁移的知识,它可以使用比教师网络更少的迭代次数,从而达到与其相同的去雾性能。
3.2.3 知识蒸馏在图像去雾中的有效性
为验证知识蒸馏在图像去雾中的有效性,将本文提出的方法与DCP 方法[2]、文献[21]方法、NLP 方法[22]、GMAN 方法[7]和EPDN 方法[6]进行对比,对比结果如图6 所示(彩色效果见《计算机工程》官网HTML 版)。其中A 组和B 组图像选自OTS 数据集,C 组和D 组图像选自HSTS 数据集。
图6 不同方法的去雾结果对比Fig.6 Dehazing results comparison among different methods
从图6 可以看出,DCP 方法对天空区域和白色区域的去雾失真严重,会将天空区域处理成深蓝色,白色物体会产生过高的对比度。文献[21]提出基于暗原色先验去雾方法,也存在天空区域去雾效果失真的问题。虽然NLP 方法去雾性能较优,但是与真实图像对比存在颜色失真的问题。基于DNN 的EPDN 方法可以生成比较自然的去雾图像,然而在雾浓度大的区域中去雾结果会留有一些薄雾。基于编码-解码结构的端到端去雾网络GMAN 通过残差模块提取雾天图像的特征信息,具有良好的去雾效果,但恢复的图像中存在较明显的网格状条纹,影响主观视觉感受。由于本文设计的基于CGAN 网络不需要生成器中特定的函数形式,也无需利用马尔科夫链反复采样、推断,因此不会存在概率问题,生成的图像也更加清晰。图6 中每幅图的下方标出了该图的去雾量化指标SSIM/PSNR,本文提出的KD-GAN 去雾方法得到的量化指标与其他方法相比取得了较高的值。表3 为在两种数据集上不同去雾方法的去雾评价指标PSNR、SSIM 的平均值。从表3可以看出,本文提出的KD-GAN 方法的PSNR 和SSIM 在OTS 和HSTS 数据集上均取得了较高的值,表明本文提出方法的去雾效果更优。
表3 不同方法的评价指标平均值对比Table 3 The average value of evaluation indexs comparison among different methods
本文提出一种基于条件生成对抗网络与知识蒸馏的单幅图像去雾方法。通过设计基于编码器-解码器结构的条件生成对抗网络,将原始有雾图像作为条件信息添加到判别器的输入端,以改进雾天图像的去雾效果。为降低生成对抗网络计算量和存储成本,通过选择L2 范数的平方来约束重型网络和轻量化网络的FSP 矩阵,使用L1 正则项对两个网络生成器输出端的特征图进行约束。在OTS 和HSTS 数据集上的实验结果表明,本文方法的峰值信噪比和结构相似性分别为0.902 0/22.418 8 和0.895 8/22.888 5,能够有效改善去雾效果,同时加快训练速度。后续将优化以过程为导向的损失函数,进一步提升学生网络的学习效率和去雾性能。