图像任务知识蒸馏调研(知识蒸馏一)

图像任务知识蒸馏调研(知识蒸馏一)

  • 图像任务知识蒸馏调研(知识蒸馏一)
    • Image Super-Resolution Using Knowledge Distillation(ACCV 2018)
      • 方法
    • Data-Free Knowledge Distillation For Image Super-Resolution(CVPR 2021)
      • 训练样本生成
      • Progressive Distillation
      • 部分实验结果
    • KD-MRI: A knowledge distillation framework for image reconstruction and image restoration in MRI workflow(MIDL2020)
      • 基于Attention的特征蒸馏
      • Imitation Loss
      • 实验结果&消融实验

图像任务知识蒸馏调研(知识蒸馏一)

最近引文研究需要,针对知识蒸馏在图像生成相关任务中的使用进行了一些调研,在这里总结成果。

Image Super-Resolution Using Knowledge Distillation(ACCV 2018)

比较早的使用知识蒸馏进行图像超分的工作,主要特点如下:

  1. 使用多层次知识蒸馏,用小网络蒸馏大网络,减少消耗(学生网络维MobileNet)
  2. 试验了不同的知识蒸馏方法,评估它们的效果

文章思路很简洁,适合初步了解知识蒸馏的应用。

方法

图像任务知识蒸馏调研(知识蒸馏一)_第1张图片
可以看到选择了三个不同层次的输出进行蒸馏,提高效率。
另外,根据以往研究,评估特征统计图的相似性要胜过直接评估特征本身的相似性,因此在进行蒸馏前,模型的输出经过了一个映射,将不同通道的信息压缩得到统计信息:
特征映射
作者选择了不同的G,得到的就是不同的方法(都是再Channel维度):

  1. 直接求平均后取幂
  2. 先取幂再求平均
  3. 取最大值
  4. 取最小值

最后结果是求平均后取二次幂效果最好。
图像任务知识蒸馏调研(知识蒸馏一)_第2张图片

Data-Free Knowledge Distillation For Image Super-Resolution(CVPR 2021)

数据蒸馏作为有效的压缩模型大小的方法已经得到了广泛应用,然而经典的蒸馏方法需要额外的训练数据用于微调模型,这在某些情况下难以做到,因此有了“data-free compression”的方式,主要的思路是根据模型来反向获得训练数据,然而这对于图像识别相关的任务有效,对于超分辨率等需要精确数据的任务来说却并不适用。文章提出的整体的模型如下所示:
图像任务知识蒸馏调研(知识蒸馏一)_第3张图片

训练样本生成

作者仍然是根据教师模型训练生成器来生成训练样本,但是和以往工作不同,在超分任务中对于样本的约束超分任务的,也就是要保持降采样一致性,损失函数如下:
损失函数1
其中 G \mathcal{G} G为生成器, z z z为随机向量, T \mathcal{T} T为教师模型, R \mathcal{R} R为降采样(相当于超分的逆操作)。

同时为了保证样本多样性,引入类似GAN的对抗损失,该函数本身和知识蒸馏相反,目的是让教师模型和学生模型对样本的结果尽量不同:
损失函数2
二者合并得到:
损失函数G

Progressive Distillation

损失函数采用经典的L1损失:
损失函数3
考虑到由于训练数据缺乏造成学生网络训练更加困难,作者选择了一种渐进式训练的方式,具体而言,就是先训练一个小网络 S 0 \mathcal{S}_0 S0,然后不断往里面插入层再训练,再插入层,直到达到需要的层数。
损失函数S

完整的训练算法如下所示,基本上可以概括为:

  1. 知识蒸馏训练学生模型
  2. 训练生成器
  3. 重复1-2
    图像任务知识蒸馏调研(知识蒸馏一)_第4张图片

部分实验结果

作者对VDSR,EDSR等模型做了蒸馏实验,在不同数据集上进行了测试,结果上看,作者的模型虽然不如在有训练数据的情况下训练的Student模型,但是显著好于Bicubic插值和使用随机Noise训练的Student模型,因此是有一定意义的:
图像任务知识蒸馏调研(知识蒸馏一)_第5张图片
同时作者进行了消融实验,验证了提出方法的有效性,以及不同 w R w_R wR的影响:
图像任务知识蒸馏调研(知识蒸馏一)_第6张图片
图像任务知识蒸馏调研(知识蒸馏一)_第7张图片
作者还窥探了一下生成器生成的样本和Teacher超分的结果(三张单通道合并成一张RGB),反正我是看不出来有什么特殊的地方。。。让我怀疑随即用一些几何色块训练也可以有好结果。。。
图像任务知识蒸馏调研(知识蒸馏一)_第8张图片

KD-MRI: A knowledge distillation framework for image reconstruction and image restoration in MRI workflow(MIDL2020)

这是(自称)第一篇将知识蒸馏用于磁共振重建的工作,作者使用DC-CNN的结构设计了Teacher和Student Network(比教师网络每个block的卷积层更少),并使用了基于attention的蒸馏损失和imitation loss(虽然这俩都不是作者原创就是了),取得了良好的效果。模型结构如下:
图像任务知识蒸馏调研(知识蒸馏一)_第9张图片

基于Attention的特征蒸馏

所谓基于Attention的特征蒸馏,实际上就是不仅仅将最终的输出作为蒸馏目标,也将网络中间层的关注点作为蒸馏目标,这在分类任务中已经被证明是有帮助的,而在image2image类型的回归任务中,这样的Attention Map更是能够直接提供关于最终生成样本的空间信息。

具体而言,对于某一层经过激活函数后的特征图 A ∈ R C × H × W A \in R^{C\times H \times W} ARC×H×W,其Attention Map可以表示为(这和本文博客第一篇文章好像啊):
Attention Map
在本文中,对Attention Map蒸馏的方法如下所示:
图像任务知识蒸馏调研(知识蒸馏一)_第10张图片
其中 Q S Q_S QS Q T Q_T QT就代表了学生和教师网络的Attention Map, j j j表示对应位置,本文中选择了每个block正中间那个(教师网络五层卷积的第三个,学生网络三层卷积的第二个)卷积层的输出,并在附录中对不同位置的效果做了验证。

Imitation Loss

为了更好地提升学生模型的效果,除了中间层的知识蒸馏,还在最后的结尾增加了Imitation损失,其实就是和Teacher网络的输出做范数损失。学生网络的输出需要分别和Ground-Truth样本和Teacher网络输出的样本做范数损失:
基于Imitation损失的总损失

其中 L i m i t = ∣ ∣ x − x r e c S ∣ ∣ L_{imit}=||x-x_{rec}^S|| Limit=xxrecS,也就是真实全采样样本和重建结果的范数损失; L i m i t = ∣ ∣ x r e c T − x r e c S ∣ ∣ L_{imit}=||x_{rec}^T-x_{rec}^S|| Limit=xrecTxrecS,注意这里 ⋅ T \cdot^T T不是转置而是教师网络的意思。

整体训练流程如下,注意两个损失的应用并不是同步的,而是分别进行的:
图像任务知识蒸馏调研(知识蒸馏一)_第11张图片

实验结果&消融实验

作者在不同数据集和不同采样率上进行了实验,使用笛卡尔降采样。使用蒸馏辅助训练的学生模型要好于不使用的。
图像任务知识蒸馏调研(知识蒸馏一)_第12张图片
在模型效率上,相比于教师模型,知识蒸馏的学生模型的参数量显著减少(114K vs. 49K),推理时间显著降低(CPU/GPU 单张 568/24ms vs. 294/16ms)。

作者还尝试了不同的蒸馏方式,以及Attention蒸馏和Imitation Loss的消融,证明了方法的有效性:
图像任务知识蒸馏调研(知识蒸馏一)_第13张图片

你可能感兴趣的:(医疗图像,计算机视觉,人工智能)