用于密集预测的通道知识蒸馏

Channel-wise Knowledge Distillation for Dense Prediction

日期：26 Nov 2020

发表：ICCV2021

作者：Changyong Shu, Yifan Liu, Jianfei Gao, Zheng Yan, Chunhua Shen

单位：Shanghai Em-Data Technology Co, The University of Adelaide, Monash University, Baidu Inc.

原文：https://arxiv.org/abs/2011.13256

开源：https://github.com/irfanICMLL/TorchDistiller/blob/main/SemSeg-distill/README.md

编者注：

（1）在github上看到一个鲜有标星的仓库(https://github.com/wzpscott/SegDistill )，是最近几个月更新的，其中提出了Channel Group Distillation，以这篇Channel-wise为基线，但是没有放论文链接，仓库贡献者包括Enze Xie

（2）在CVPR2022上也有一篇语义分割知识蒸馏的工作(https://github.com/winycg/CIRKD)，是把这篇CWD作为SOTA

（3）有另外个github仓库开源，从这个issue上看，应该也是来自作者团队：https://github.com/drilistbox/CWD/issues/1

摘要：知识蒸馏（KD）已被证明是训练紧凑预测模型的一种简单有效的工具。轻量级学生网络通过从大型教师网络迁移的额外监督进行训练。用于密集预测任务的大多数先前KD变体通常通过归一化每个空间位置上的激活值和最小化逐点和/或逐对差异，在空间域中对齐来自学生和教师网络的激活图。与之前的方法不同，这里我们建议归一化每个通道的激活图以获得软概率图。通过简单地最小化两个网络的通道概率图之间的Kullback–Leibler（KL）散度，蒸馏过程更加关注每个通道的最显著区域，这些区域对于密集预测任务很有价值。我们在一些密集预测任务上进行了实验，包括语义分割和对象检测。实验表明，我们提出的方法大大优于现有的蒸馏方法，并且在训练过程中需要更少的计算成本。特别是，我们在COCO数据集上的mAP中将RetinaNet检测器（ResNet50主干）提高了3.4%，在城市景观数据集上，在mIoU中将PSPNet（ResNet18主干）提高5.81%。代码可在https://github.com/irfanICMLL/TorchDistiller/blob/main/SemSeg-distill/README.md

Knowledge distillation (KD) has been proven to be a simple and effective tool for training compact models. Almost all KD variants for dense prediction tasks align the student and teacher networks' feature maps in the spatial domain, typically by minimizing point-wise and/or pair-wise discrepancy. Observing that in semantic segmentation, some layers' feature activations of each channel tend to encode saliency of scene categories (analogue to class activation mapping), we propose to align features channel-wise between the student and teacher networks. To this end, we first transform the feature map of each channel into a probabilty map using softmax normalization, and then minimize the Kullback-Leibler (KL) divergence of the corresponding channels of the two networks. By doing so, our method focuses on mimicking the soft distributions of channels between networks. In particular, the KL divergence enables learning to pay more attention to the most salient regions of the channel-wise maps, presumably corresponding to the most useful signals for semantic segmentation. Experiments demonstrate that our channel-wise distillation outperforms almost all existing spatial distillation methods for semantic segmentation considerably, and requires less computational cost during training. We consistently achieve superior performance on three benchmarks with various network structures. 

1    引言

密集预测任务是计算机视觉中的一组基本任务，包括语义分割[48，6]和目标检测[21，30]。这些任务需要学习像素级复杂场景理解的强特征表示。因此，最先进的模型通常需要较高的计算成本，这使得它们难以部署到移动设备。因此，为密集预测任务设计的紧凑网络引起了广泛关注。此外，在以前的工作中，使用知识蒸馏（KD）对有效训练轻量级网络进行了研究。紧凑型网络在大型教师网络的监督下进行训练，可以获得更好的性能。提出并研究了开创性工作[16，2]，主要用于图像分类任务。

密集预测任务是逐像素预测问题，比图像级分类更具挑战性。先前的研究[25，20]发现，将分类中的KD方法[16，2]直接迁移到语义分割可能不会产生令人满意的结果。在教师和学生网络之间严格对齐逐点分类分数或特征图可能会强制执行过于严格的约束并导致次优解。

最近的工作[25、24、18]关注加强不同空间位置之间的相关性。如图2（a）所示，每个空间位置上的激活值（原文注：本文中，激活值包含最终的logits和中间的特征图）被归一化。然后，通过聚集不同空间位置的子集来执行一些特定于任务的关系，例如成对关系[25，35]和类间关系[18]。这种方法在捕获空间结构信息和提高学生网络的性能方面可以比逐点对齐更好地工作。然而，激活图中的每个空间位置对知识迁移的贡献相等，这可能会从教师网络中带来冗余信息。

在这项工作中，我们提出了一种新的基于通道的知识提取方法，通过对密集预测任务的每个通道中的激活图进行归一化，如图2（b）所示。然后，我们最小化教师和学生网络的归一化通道激活图之间的不对称Kullback–Leibler（KL）散度，该散度被转换为每个通道的分布。我们在图2（c）中展示了通道分布的示例。每个通道的激活倾向于编码场景类别的显著性。对于每个通道，引导学生网络更加注意模拟具有显著激活值的区域，从而在密集预测任务中实现更精确的定位。例如，在对象检测中，学生网络更加关注学习前景对象的激活。

图2。左上角：空间知识蒸馏的工作对齐空间域的特征图。右上角：我们的通道分布蒸馏对齐学生和老师的特征图的每个通道，通过最小化KL散度。下面一栏显示每个通道的激活值倾向于编码场景类别的显著信息。

最近的一些作品利用了通道中包含的知识。通道蒸馏[50]提出将每个通道中的激活转换为一个聚集标量，这可能有助于图像级分类，但空间聚集会丢失所有空间信息，因此不适合密集预测。其他工作，如MGD[41]、通道交换[33]和CSC[26]表明了通道信息的重要性。MGD将教师通道与学生通道相匹配，并将其作为作业问题解决。通道交换[33]使用融合模块在各种模式的子网络之间动态交换通道。

我们表明，每个通道的简单归一化操作可以大幅度提高基线空间蒸馏。所提出的通道式蒸馏方法简单，易于应用于各种任务和网络结构。我们将主要贡献总结如下。

•    与现有的空间蒸馏方法不同，我们为密集预测任务提出了一种新的按通道蒸馏范式。我们的方法简单而有效。

•    在语义分割和对象检测方面，提出的逐通道蒸馏方法明显优于最先进的KD方法。

•    在语义分割和对象检测任务方面，我们对四个具有不同网络结构的基准数据集进行了一致的改进，证明了我们的方法是通用的。鉴于其简单性和有效性，我们认为我们的方法可以作为密集预测任务的强基线KD方法。

2    相关工作

大多数关于知识提取的工作集中于分类任务[11、12、16、27、36、38、45]。我们在这里的工作旨在研究用于密集预测的高效蒸馏方法，而不是像在分类中那样简单地应用逐像素蒸馏。

语义分割的知识蒸馏。

在[35]中，构建了局部相似性图，以最小化教师和学生网络之间分割边界信息的差异，其中中心像素和8邻域像素之间的欧几里得距离被用作知识传递。Liu等人[24，25]提出了两种捕获像素间结构化信息的方法，包括像素间的成对相似性和鉴别器捕获的整体相关性。[34]中的工作集中于具有相同标签的像素之间的类内特征变化，其中构造每个像素的特征与其对应的类间原型之间的余弦距离集，以传递结构知识。他等人[14]使用特征适配器来缓解教师和学生网络之间的特征不匹配。

用于对象检测的知识蒸馏。

许多方法发现，在目标检测的蒸馏中，区分前景和背景区域是重要的。MIMIC[20]通过L2损失迫使学生网络的RPN内的特征图与教师网络的特征图相似，并发现直接应用像素损失可能会损害对象检测的性能。Wang等人[32]提出提取对象锚位置附近的细粒度特征。Zhang和Ma[43]在生成掩模时注意区分前景和背景，取得了令人满意的结果。相反，我们柔和地对齐通道激活，以区分前景和背景区域。

通道的知识。

最近的几个工作[50]也关注每个通道中包含的知识。Zhou等人计算每个通道中激活的平均值，并在分类中对齐每个通道的加权差。CSC[26]计算用于传输知识的所有空间位置和所有通道之间的成对关系。通道交换[33]提出，每个通道中包含的信息是通用的，可以跨不同的模式共享。

3    我们的方法

我们首先回顾了文献中相关的空间知识提取方法。

3.1    空间蒸馏

现有的KD方法通常采用逐点对齐或在空间位置之间对齐结构化信息，其可以表示为：

    （1）

在这里，任务损失仍然使用，是真值标签。例如，交叉熵损失通常用于语义分割。通过稍微滥用符号，这里和分别表示学生和教师网络的logits或内部激活。这里，是平衡损失项的超参数。下标和表示教师和学生网络。我们在表1中列出了代表性的空间蒸馏方法。

这些方法的简要概述如下。注意力迁移（AT）[42]使用注意力掩码将特征映射压缩到单个通道中进行蒸馏。逐像素损失[17]直接对齐逐点类概率。局部亲和度计算为通过中心像素与其8个相邻像素之间的距离[35]。成对亲和度[25，14，24]计算像素对之间的相似性。传递结构知识计算每个像素的特征与其对应的分类原型之间的相似度[34]。[25，24]中的整体损失使用对抗方案来对齐两个网络的特征图之间的高阶关系。注意，最后四项考虑像素之间的相关性。表1所示的现有KD方法都是空间蒸馏方法。所有这些方法都将某个空间位置的N通道激活值视为要操作的特征向量。

3.2    通道式蒸馏

为了更好地利用每个通道中的知识，我们建议在教师和学生网络之间柔和地调整相应通道的激活。为此，我们首先将通道的激活转换为概率分布，以便我们可以使用概率距离度量（如KL散度）来测量差异。如图2（c）所示，不同通道的激活倾向于编码输入图像场景类别的显著性。此外，训练好的语义分割教师网络显示了每个通道的清晰类别特定掩码的激活图，如图1右侧所示。在这里，我们提出了一种新的按通道蒸馏范式，以指导学生从训练好的教师那里学习知识。

让我们将教师和学生网络表示为和，来自和的激活图分别为和。通道式蒸馏损失可以用一般形式表示：

    （2）

在我们的情况下，用于将激活值转换为概率分布，如下所示：

    （3）

其中通道索引；索引某个通道的空间位置。是超参数（温度）。如果我们使用更大的，概率会变得更软，这意味着我们关注每个通道的更宽空间区域。通过应用softmax归一化，我们消除了大型网络和紧凑网络之间的震级尺度的影响。如[31]所述，这种归一化有助于KD。如果教师和学生之间的通道数量不匹配，则采用1×1卷积层对学生网络的通道数量进行上采样。评估教师网络和学生网络的通道分布之间的差异。我们使用KL散度：

    （4）

KL散度是一个非对称度量。从等式（4）中，我们可以看出，如果较大，则应与一样大，以最小化KL散度。另一方面，如果非常小，KL散度不太注意最小化。因此，学生网络倾向于在前景显著性中产生类似的激活分布，而教师网络的背景区域相对应的激活对学习的影响较小。我们假设KL的这种不对称性有利于密集预测任务的KD学习。

4    实验

在本节中，我们首先描述实现细节和实验设置。然后，我们将我们的逐通道蒸馏方法与其他最先进的蒸馏方法进行比较，并对语义分割进行消融研究。最后，我们通过各种基准和学生网络结构展示了语义分割和对象检测的一致改进。

4.1    实验设置

数据集。

这里使用了三个公共语义分割基准，即Cityscapes[8]、ADE20K[49]和Pascal VOC[10]。我们还将所提出的蒸馏方法应用于MS-COCO 2017[23]上的目标检测，这是一个包含80个类别的120k多幅图像的大规模数据集。

城市景观数据集用于语义城市场景理解。它包含5000幅精细注释图像，分别有2975/500/1525幅图像用于训练/验证/测试，其中提供了30个常用类，19个类用于评估和测试。每个图像的大小为2048×1024像素。他们来自50个不同的城市。我们的实验中不使用粗标记数据。

Pascal VOC数据集包含1464/1449/1456幅图像，用于训练/验证/测试。它包含20个前景对象类和一个额外的背景类。此外，该数据集还通过额外的粗标记进行了扩展，其中有10582幅图像用于训练。训练分割用于训练，并在21个类的验证集上测量最终性能。

ADE20K数据集涵盖150类不同场景。它包含用于训练、验证和测试的20K/2K/3K图像。在我们的实验中，我们报告了验证集上的分割精度。

评估指标。    

为了评估我们提出的通道分布蒸馏方法在语义分割方面的性能和效率，继之前的工作[18，24]之后，我们在单尺度设置下的所有实验中通过交并比（mIoU）来测试每个策略。每秒浮点运算（FLOPs）是在512×1024像素的固定输入大小下计算的。此外，还列出了Pascal VOC和ADE20K的平均类精度（mAcc）。为了评估目标检测的性能，我们报告了平均平均精度（mAP）、推理速度（FPS）和模型大小（参数），如[43]所述。

实现细节。    

对于语义分割，教师网络是PSPNet，ResNet101作为所有实验的骨干（PSPNet-R101）。学生网络，我们采用了几种不同的体系结构，包括以ResNet18和MobileNetV2为主干的DeepLab[44]、PSPNet[48]，以验证我们方法的有效性。

在消融研究中，我们分析了基于具有ResNet18主干（PSPNet-R18）的PSPNet的方法的有效性。除非另有说明，学生网络的每个训练图像被随机裁剪成512×512像素。批量大小设置为8，训练步骤的数量为40K。对于所有实验，我们将温度参数=4，logits图的损失重量=3，特征图的损失系数=50。对于目标检测，我们采用与[43]中相同的教师和学生网络以及训练设置。

4.2    与最近的知识蒸馏方法的比较

为了验证我们提出的逐通道蒸馏的有效性，我们将我们的方法与下面列出的当前蒸馏方法进行了比较：

•    注意力迁移（AT）[42]：Sergey等人计算每个空间位置处所有通道的总和，以获得单通道注意力图。L2用于最小化注意力图之间的差异。

•    局部相似性（LOCAL）[35]：对于每个像素，构建局部相似性图，该图考虑了自身与其8个相邻像素之间的相关性。L2用于最小化局部亲和图之间的差异。

•    逐像素蒸馏（PI）[25、24、34、7]：KL散度用于对齐两个网络中每个空间位置的分布。

•    成对蒸馏（PA）[25、14、24]：考虑所有像素对之间的相关性。

•    类内特征变化蒸馏（IFVD）[34]：每个像素的特征与其对应的类间原型之间的相似性集合被视为类内特征差异，以传递结构知识。

•    整体蒸馏（HO）[25、24、34]：特征图的整体嵌入由鉴别器计算，该鉴别器用于最小化高阶关系之间的差异。

我们将所有这些流行的蒸馏方法应用于内部特征映射和最终logits映射。传统的交叉熵损失应用于所有实验。表2报告了空间蒸馏方法的计算复杂性和性能。

给定（）大小的输入特征图（logits图），其中是特征图（logits图）的形状。是通道数，是类数。

如表2所示，所有蒸馏方法都可以提高学生网络的性能。我们的通道蒸馏方法优于所有空间蒸馏方法。我们的方法比最佳空间蒸馏方法（AT）高出2.5%。此外，我们的方法更高效，因为它在训练阶段比其他方法需要更少的计算成本。

此外，我们在表3中列出了我们方法的详细类IoU和两种最新的方法，PA[25]和IFVD[18]。这些方法提出在语义分割中传递结构信息。我们的方法显著提高了交通灯、地形、墙壁、卡车、公共汽车和火车等多个对象的分类精度，表明通道分布可以很好地传递结构知识。

4.3    消融研究

在本节中，我们展示了逐通道蒸馏的有效性，并讨论了语义分割中超参数的选择。基线学生模型为PSPNet-R18，教师模型为PSPNet-R101。所有结果均在城市景观验证集上进行评估。

通道式蒸馏的有效性。

归一化通道概率图和非对称KL散度在我们的蒸馏方法中起着重要作用。表4中我们使用四种不同的变体进行实验，以显示提出的方法的有效性。

所有蒸馏方法都使用相同的激活图作为输入。我们使用与第4.1节所述相同的训练方案。

“PI”表示像素级知识蒸馏，它归一化每个空间位置的激活。L2 w/o NORM表示我们直接最小化两个网络的特征图之间的差异，它同等地考虑了所有通道中所有位置的差异。Bhat是Bhattacharyya距离[3]，这是一种对称分布测量。它对齐每个通道中的差异。

从表4中，我们可以看出，度量归一化通道差异的非对称KL散度实现了最佳性能。注意，由于KL散度是不对称的，学生和教师的输入不能交换。我们通过改变KL散度中输入的顺序进行实验，训练不会收敛。

温度参数和损失权重的影响。我们通过在不同的损失权重下调整温度参数来进行实验，以改变通道概率图。实验在logits图上进行。结果如图3所示。

所有结果均为三次运行的平均值。损失权重设置为1、2、3和∈ [1, 5]. 如果我们增加，分布趋于更软。

从图中，我们可以看出，更柔和的概率图可能有助于知识蒸馏。此外，在一定范围内，性能稳定。如果设置得很小，性能似乎会下降。在这种情况下，该方法仅关注有限的显著像素。当=4和=3时，我们在城市景观验证集上使用PSPNet18获得最佳性能。

4.4    语义分割

我们证明，在各种学生网络下，我们提出的逐通道蒸馏方法可以与以前的语义分割蒸馏方法相结合，即用于分割/密集预测的结构知识蒸馏（SKDS[24]和SKDD[25]）和类内特征变化蒸馏（IFVD[34]）。

我们在logits映射（Ours-logits）和特征映射（Ours-feature）上使用所提出的逐通道蒸馏。logits图上的逐像素蒸馏（PI）和整体蒸馏（HO）也包括在前面的方法[25，18]中。

我们首先在城市景观数据集上评估我们的方法的性能。使用具有不同编码器和解码器的各种学生网络来验证我们方法的有效性。编码器包括ResNet18（使用或不使用在ImageNet上预训练的权重初始化，以及ResNet18的通道减半变体[13]），解码器包括PSPHead[48]和ASPPHead[6]。表5显示了城市景观的结果。Pascal VOC[10]和ADE20K[49]的实验结果见补充材料。

在五个学生网络和三个基准上，我们的方法优于SKD和IFVD，这进一步表明逐通道蒸馏对于语义分割是有效的。

对于与教师具有相同结构类型的学生，即PSPNet-R18（0.5）、PSPNet-R18和PSPNet-R18，这些改进更为显著。对于与教师有不同架构类型的学生，即Deeplab-R18*（0.5）和Deeplab-R18，与SKDS和IFVD相比，我们的方法实现了一致的改进。因此，我们的方法适用于不同的教师和学生网络。

与具有较大容量的学生（PSPNet-R18）(75.90%)相比，具有紧凑模型容量的学生网络（PSPNet-R18（0.5））显示出较差的蒸馏性能（68.57%）。这可能是由于与教师网络相比，小型网络的能力有限，无法充分吸收当前任务的知识。对于PSPNet-R18，通过在ImageNet上训练的权重初始化的学生获得了最佳的蒸馏性能（从70.09%提高到75.90%），进一步证明了初始化良好的参数有助于蒸馏。因此，更好的学生会导致更好的蒸馏性能，但由于师生网络之间的差距较小，这种改进并不显著。(Thus, the better student lead to better distillation performance, but the improvement is less significant as the gap between the teacher and student network is smaller.)

4.5    物体检测

我们还将我们的逐通道蒸馏方法应用于目标检测任务。实验在MS COCO2017[23]上进行。

使用不同范式下的各种学生网络，即两阶段基于锚的方法（Faster RCNN[28]）、一阶段基于锚方法（RetinaNet[21]）和无锚方法（RepPoints[37]）来验证我们方法的有效性。为了进行公平的比较，我们使用与[43]中相同的超参数对同一位教师进行了实验。

唯一的修改是将特征对齐更改为我们的逐通道蒸馏。结果如表6所示。从表中，我们可以看出，我们的方法在强基线学生网络上实现了一致的改进（约3.4%的mAP）。与之前最先进的蒸馏方法[43]相比，我们的简单通道式蒸馏性能更好，尤其是无锚方法。我们将报告点提高了3.4%，而Zhang等人将报告点提升了2%。此外，我们可以看到，提出的蒸馏方法可以更显著地改善AP75。

5    结论

在本文中，我们提出了一种新的用于密集预测任务的逐通道蒸馏。与以前的空间蒸馏方法不同，我们将每个通道的激活归一化为概率图。然后，应用不对称KL发散来最小化教师和学生网络之间的差异。实验结果表明，在四个具有不同网络主干的公共基准数据集上，所提出的蒸馏方法在语义分割和对象检测方面都优于现有的蒸馏方法。

此外，我们的烧蚀实验证明了我们的通道式蒸馏的效率和有效性，它可以进一步补充空间蒸馏方法。我们希望，所提出的简单有效的蒸馏方法可以作为有效训练紧凑网络的强大基线，用于许多其他密集预测任务，包括实例分割、深度估计和全景分割。

附录

A    Pascal VOC和ADE20K的结果

为了进一步证明所提出的通道分布蒸馏的有效性，我们仅在特征图上使用所提出的CD作为我们对Pascal VOC和ADE20K的最终结果。实验结果如表7和表8所示。使用具有不同编码器和解码器的多学生网络变体来验证我们方法的有效性。这里，编码器包括ResNet18和MobileNetV2，解码器包括PSP头和ASPP头。

Pascal VOC

我们在Pascal VOC数据集上评估了我们的方法的性能。蒸馏结果如表7所示。我们提出的CD相比不蒸馏的情况提高了PSPNet-R18 3.83%，优于SKDS和IFVD 1.51%和1.21%。在具有不同编码器和解码器的其他学生网络上实现了一致的改进。使用我们的方法在PSPNet-MBV2上的增益为3.55%，超过SKDS和IFVD 1.98%和1.20%。至于Deeplab-R18，我们的CD将学生从66.81%提高到69.97%，分别比SKDS和IFVD高出1.84%和1.55%。此外，通过我们的蒸馏，DeepLab-MBV2的性能从50.80%提高到54.62%，分别优于SKDS和IFVD 2.51%和1.23%。

ADE20K

我们还在ADE20K数据集上评估了我们的方法，以进一步证明CD比其他结构知识提取方法更有效。结果如表8所示。我们提出的CD相比不蒸馏的情况将PSPNetR18提高了3.83%，并在几个方面优于SKDS和IFVD 1.51%和1.21%。在使用不同编码器和解码器的其他学生身上，也始终实现了显著的性能提升。对于PSPNetMBV2，我们的方法实现了27.97%的优异性能，超过了学生、SKDS和IFVD 4.82%、3.18%和2.64%。使用我们的CD在Deeplab-R18上的增益为2.48%，比SKDS和IFVD分别高出1.85%和0.84%。最后，使用我们的通道式蒸馏的Deeplab-MBV2的性能从24.98%提高到29.18%，分别优于SKDS和IFVD 3.08%和1.93%。

B    更多可视化结果

我们在图6中列出了可视化结果，直观地证明了通道分布蒸馏方法（CD）优于空间蒸馏策略（注意力迁移）。此外，为了评估所提出的通道分布蒸馏的有效性，我们在图7和图8中，分别通过注意力迁移（AT）和通道分布蒸馏，可视化了三种范式下的学生网络的通道分布，即原始网络。