《Structured Knowledge Distillation for Dense Prediction》论文笔记

代码地址：structure_knowledge_distillation

1. 概述

导读：这篇文章针对的是密集预测的网络场景（如语义分割），在之前的一些工作中对于这一类网络的蒸馏时照搬分类任务中那种逐像素点的蒸馏方式（相当于是对每个像素的信息分别进行蒸馏），文章指出这样的产生的结果并不是最优的（这样策略会忽视特征图里面的结构信息，像素信息之间是存在关联的），因而这篇文章提出了适应密集预测网络的蒸馏策略：1）pair-weise蒸馏：通过构建静态图（受pair-wise的马尔可夫随机场启发，增强特征图中sptial上的的相关性，使得可以学习到结构性信息）蒸馏对应的相似性；2）holistic蒸馏：使用对抗训练去蒸馏学生网络和教师网络的输出信息（这里使用了更加高维度的信息，目的是使得判别器无法判断信息的来源）。文章的方法在语义分割，深度估计与目标检测上进行了实验，其结果也显示了其有效性。

文章从传统的蒸馏方法开始进行分析，由于传统的蒸馏方法是逐像素点的蒸馏方式，对特征中的结构性信息并没有很好提取，对此文章针对性的给出两种蒸馏的策略：

• 1）pair-wise方式的蒸馏：文章使用pair-wise的马尔科夫随机场框架来增强空间labelling的连续性，目标是对齐简单网络（student）和复杂网络（teacher）中学到的pair-wise特征，从而使得学生网络能够学习到更多的结构信息；
• 2）holistic蒸馏：这里并不将知识迁移的维度限定在pair-wise与pixel-wise上，而是使用对抗训练的形式监督学生与教师网络的输出，使其在更高的维度上进行近似逼近。判别器考虑的是网络输入图像（作为条件输入）与网络输出组成的holistic embedding，使得学生网络生成的结果不断近似教师网络。

使用文章的方法进行蒸馏，其在相应的baseline上得到的性能比较见下图所示：

2. 方法设计

2.1 蒸馏的整体结构

对于分割任务其流程是：对于个3通道的输入图像$I\in R^{W\ast H\ast 3}$，它在经过卷积网络特征抽取之后，得到特征图$F\in R^{W^{‘}\ast H^{‘}\ast N}$的特征图（论文代码中给出的stride=8），之后对其使用分类得到分类类别为C的结果$Q\in R^{W^{‘}\ast H^{‘}\ast C}$，之后将其上采样与原始输入图像尺寸保持一致。因而对于像分割这类的密集预测问题，文章设计了图2的蒸馏结构：

在上图中总共设计了3种类型的蒸馏损失：pair-wise的蒸馏损失，pixel-wise的蒸馏损失，Wasserstein距离损失（学生网络还有分割交叉墒损失）。

2.2 Pixel-wise蒸馏

文中使用标记$S$代表学生网络，$T$代表教师网络。对于分割部分特征图$Q$其使用的是原始的蒸馏方式，使用KL散度计算差异，因而这部分的损失函数描述为：
$$L_{pi}(S)=\frac{1}{W^{‘}\ast H^{‘}}\sum _{i\in R}KL(p_i^s,p_i^t)$$
其中，$p_i^s,p_i^t$代表来自学生网络与教师网络的概率值， R = { 1 , 2 , … , W ‘ ∗ H ‘ } R=\{1,2,\dots,W^{‘}*H^{‘}\} R={ 1,2,…