Distilling image dehazing with heterogeneous task imitation

1、摘要

最先进的深度除雾模型在训练中往往是困难的。知识蒸馏使得利用教师网络训练学生网络成为可能。然而，大多数的知识蒸馏方法都是应用于图像分类、语义分割或者目标检测的，很少有研究将知识蒸馏应用于不同图像恢复任务之间的知识转移。本文提出了一种利用异质任务模拟提取图像去雾知识的网络，在本文所提出的模型中，teacher网络是一个自编码器，用于图像重建，去雾网络在teacher网络的辅助下，采用过程导向机制进行学习。student网络学习teacher网络以进行图像重建任务。此外，本文为student网络设计了一个空间加权的通道注意力残差块，以使其自适应地关注通道级别的特征，并更加关注浓雾区域的重建恢复。

2、引言

2.1 知识蒸馏去雾的难点包括：

1）什么样的网络在什么任务下可以有效的为图像去雾任务提供有效的辅助与指导？现有的知识蒸馏方法大多是使用Teacher和Student网络来处理相同的任务，很少有人研究不同任务的两个网络之间的知识转移，异构任务是否有助于深度去雾模型的训练？

2）教师模型应该如何去辅助学生模型，目前大多数的知识蒸馏方法侧重于结果导向的学习，而忽视了过程导向的学习。如何利用学习过程中的暗知识去辅助student模型去训练是目前知识蒸馏所需要解决的一个问题。

3）如何衡量Teacher和Stude之间的相似性?在基于知识蒸馏的图像分类中，可以利用类别分布来衡量Teacher和Student网络之间的相似度。然而，对于图像恢复任务，是没有类别分布可以用来衡量相似度的。

针对以上三个问题，解决分别如下：首先，使用可以重建图像的自编码器作为教师网络以对去雾模型进行监督；其次，采用面向过程的知识转移学习机制，对教师网络的中间特征进行监督，利用特征相似性来使得学生网络与教师网络的重建过程是相似的。面向过程的监督充分利用了暗知识，可以获得更好的图像去雾效果。最后，针对第三个问题，利用图像保真度损失函数、感知损失函数以及Teacher和Student在中间过程中生成的特征图的差异来解决。

2.2 本文主要贡献包括：

1）利用异构知识蒸馏构建了一种去雾网络。

2）提出过程导向学习机制。

3）提出了 spatial-weighted channel-attention residual block 与 a haze density aware imitation loss。

3、提出的模型

知识蒸馏去雾网络（knowledge distilling dehazing network，KDDN）的整体架构由Teacher/Student网络组成。具体来说，Teacher的目标是提供 clear image 的中间特征表示，Student的目标是通过拟合Teacher网络中 clear image 的特征，将clear image 从 haze image 中恢复出来。

3.1 网络架构

3.1.1 Teacher Network

如上图所示，由下采样层、残差块、上采样层组成。其中，下采样：跨步卷积；基本残差结构：卷积+RELU+卷积+ADD操作；上采样：插值+卷积。通道数设置为64，所有网络均不使用BN层，残差块的个数为6。

3.1.2 Student Network

总体结构与 Teacher Network 基本相同，不同点在于基本模型的设置，Student 网络使用了多个RIR模块。RIR模块一般由多个残差块或者SE模块组成，但是，SE模块使用全局池化来学习每个通道的权重，忽略了图像的不同位置雾霾浓度时不一致的这一事实。为了解决这一问题，本文使用了空间加权残差注意力模块（spatially weighted residual attention block， SWRCAB)，如图所示，SWRCAB首先通过卷积与sigmoid操作学习输入特征的空间权重，然后使用得到的权重对经过两次卷积操作后的输入特征（记为F）进行加权，加权后的特征经过全局池化、全连接层、sigmoid得到通道的权重，然后使用使用通道权重对F进行加权。通道加权后的特征与输入特征相加，得到最终的特征。

3.2 损失函数

3.2.1 Teacher网络

使用L1损失

3.2.2 Student网络

总损失为：重建损失+感知损失+表征模仿损失（representation mimicking loss）

重建损失可以表示为：

感知损失可以表示为：

表征模仿损失可以表示为：

其中，T表示教师网络，m表示第m层特征，S表示学生网络，n表示第n层特征。

考虑到浓度的位置退化较为严重，网络的浅层难以对这些位置进行恢复，所以本文对表征模仿损失进行了加权，即加大浓度部分的损失的权重。加权后的损失可以表示为：

作者给出了两种加权方式，一种根据透射率图（这是最直接反应雾浓度的变量），另一种是根据有雾图与无雾图像素值之差（因为有时透射率是一个未知变量），norm表示归一化，使用了极大极小值归一化方式：

4. 实验结果

4.1 训练细节

Teacher网络与Student网络是分开训练的，首先训练30个epoch 的Teacher网络，然后训练60个epoch的Student网络。

4.2 消融研究

4.2.1 不同损失

4.2.2 不同程度的Teacher网络

random表示随机初始化，w/o表示没有监督，wM表示中程度监督，wH表示高程度监督。

知识蒸馏去雾：Distilling image dehazing with heterogeneous task imitation