大道至简，无痛涨点AttnFD | 注意力引导特征蒸馏用MSE Loss即可成就Cityscapes巅峰mIoU精度

集智书童公众号

发布于 2024-03-14 10:59:21

1.1K0

发布于 2024-03-14 10:59:21

文章被收录于专栏：集智书童

与目前常用于将知识从教师模型传递给学生模型的复杂方法相比，所提出的方法展示了利用细化特征图来转移注意力的简单而强大的方法的有效性。该方法在提炼丰富信息方面被证明是有效的，并且在作为密集预测任务的语义分割中，其表现超越了现有方法。所提出的注意力引导特征蒸馏（AttnFD）方法采用了卷积块注意力模块（CBAM），它通过考虑通道特异性和空间信息内容来细化特征图。只使用教师和学生细化特征图之间的均方误差（MSE）损失函数，AttnFD在语义分割方面表现出色，在PascalVoc 2012和Cityscapes数据集上，平均交并比（mIoU）达到了最新水平。代码：https://github.com/AmirMansurian/AttnFD

1 Introduction

语义分割在计算机视觉中是一项非常重要且具有挑战性的任务。它已成为各种应用中不可或缺的组成部分，例如自动驾驶、视频监控和场景解析。其目标是通过为图像中的每个像素分配一个特定的类别标签来实现密集预测。通过使用深度神经网络，特别是全卷积网络（FCN），语义分割已经取得了显著的进展。其他方法通过在FCN的基础上进行构建，持续提高了分割的准确性。它们通过采用如下策略来实现这一点：设计更深的架构以增加FCN的容量，融入更强的基础网络，以及分层处理图像上下文。增加复杂性在提高语义分割的准确性方面是有效的，然而在资源受限的环境中，如移动和边缘设备，这已成为一个日益凸显的问题。

近年来，许多研究致力于设计具有较低计算成本的轻量级模型，以适应现实世界中的应用。例如，它们通过减少卷积层的数量使复杂的 Backbone 网络变得更轻，或者完全用更简单的模型替换 Backbone 网络来提高效率。在这些方法中，知识蒸馏（Knowledge Distillation，KD）已被证明是优化深度神经网络中准确性与效率之间平衡的有效策略。这项技术从较大的（教师）网络中提炼出有用的信息，如像素间的关联，并利用这些知识来监督较轻的（学生）网络的训练。尽管像素级的知识蒸馏对图像分类有益，但由于其在捕捉像素间上下文相关性方面的局限性，它在提高语义分割任务的性能上面临挑战。

随着研究的进展，已经转向基于特征的知识蒸馏，并在教师网络和学生网络之间对齐中间特征图。为了实现这一点，许多方法提出了复杂的损失函数来增强知识蒸馏，因为简单地用距离度量复制特征图存在局限性。尽管这些方法有效，但最近的研究表明，在保留基本损失函数的同时，通过新型模块转换学生特征图可以导致网络结构更简单且性能得到提升。

注意力机制旨在模仿人类观察视觉场景的方式。实际上，人类往往不是一次性处理整幅图像，而是倾向于选择含有重要信息的区域给予更多关注（通过忽略图像的其他部分）。鉴于注意力在聚合长距离上下文语义信息方面的能力，将其融入到基于特征的知识蒸馏中，对于产生重大影响颇具潜力。尽管这一方法具有潜在的效用，但目前在很大程度上还未被探索。

与之前的工作不同，前者要么定义复杂的损失来考虑成对关系，要么依赖于原始特征，本研究采用CBAM中的注意力机制。这种机制结合了通道和空间信息来生成精炼的特征，然后从教师网络传递到学生网络。图1展示了原始特征与精炼特征之间的区别。如图所示，精炼特征突出了图像的重要区域并减少了背景噪声，这使得它成为具有显著潜力用于知识蒸馏的强候选者。这是因为它迫使学生网络模仿教师网络强调的重要区域。

本工作主要贡献的总结包括：

提出了一种简单而有效的基于注意力的特征蒸馏方法用于语义分割。通过利用教师网络和学生网络的原始特征图，通道和空间注意力被用于生成用于蒸馏的细化特征图，通过采用这个模块，引入了一种新颖的知识蒸馏方法。
在知识蒸馏（KD）用于语义分割方面超越现有方法。所提出的注意力引导特征蒸馏方法显著提升了紧凑模型的最新性能，这一改进在广泛使用的基准数据集上得到了验证。

3 Proposed Method

在本节中，提出了所建议的注意力引导特征蒸馏的公式。所提出方法的总览展示在图2中。

令

A\in\mathbb{R}^{c\times w\times h}

表示从学生网络或教师网络中获得的一个中间特征图，其空间维度为

h\times w

，通道维度为

。在网络的不同阶段，两个注意力模块聚集中间特征的空域和通道描述符

M_{S}(A)\in\mathbb{R}^{c\times h\times w}

和

M_{C}(A)\in\mathbb{R}^{c\times 1\times 1}

。这些空域和通道注意力描述符随后与原始特征图

相乘，以创建新的富含上下文信息的特征图

A^{{}^{\prime}}

和

A^{{}^{\prime\prime}}

，将类间和空间信息引入原始特征图中。特征细化的整体公式如下:

A^{{}^{\prime}} =M_{C}(A)\otimes A \tag{1}

A^{{}^{\prime\prime}} =M_{S}(C)\otimes C \tag{2}

其中

\otimes

表示逐元素的乘法。在乘法过程中，空间注意力图沿通道广播，而通道注意力图沿空间维度广播。通道和空间注意力模块的方法分别可见于图3和图4。

Channel Attention Module

通道注意力模块（CAM），如图3所示，通过将最大池化和平均池化操作应用于中间特征图

来聚合空间信息。这些操作产生上下文描述符，然后通过多层感知机进行处理以生成通道注意力图

M_{C}(A)

。这张图突出了图像中有意义的区域，同时模糊与分割任务不相关的区域，比如背景。中间特征图

的通道感知描述符定义如下：

M_{C}(A)=\sigma(W_{1}(W_{0}(A^{C}_{\text{avg}})))+W_{1}(W_{0}(A^{C}_{\text{max }})) \tag{3}

其中，

A^{C}_{\text{avg}}\in\mathbb{R}^{c\times 1\times 1}

和

A^{C}_{\text{max}}\in\mathbb{R}^{c\times 1\times 1}

是通过对中间特征图

A\in\mathbb{R}^{c\times h\times w}

应用平均池化和最大池化算子生成的特征图。

W_{1}

和

W_{0}

是两个池化特征图之间共享的MLP的权重。

W_{0}

后面跟着一个ReLU激活函数，而

\sigma

表示sigmoid函数。

Spatial Attention Module

空间注意力模块（SAM），如图4所示，其操作与CAM非常相似。中间特征图

A\in\mathbb{R}^{c\times h\times w}

的空间信息通过使用最大池化和平均池化操作来聚合，生成两个不同的空间上下文描述符

A^{S}_{\text{avg}}\in\mathbb{R}^{1\times h\times w}

和

A^{S}_{\text{max}}\in\mathbb{R}^{1\times h\times w}

。

M_{S}(A)=\sigma(f^{7\times 7}([A^{S}_{\text{avg}};A^{S}_{\text{max}}])) \tag{4}

在哪里，

f^{7\times 7}

代表一个

7\times 7

的卷积核。使用新获取的特征，所提出的损失计算为：

L_{Attn}=\frac{1}{N}\sum_{i=1}^{N}\left\|\frac{A^{{}^{\prime\prime}}_{\mathcal{ S}j}}{\parallel A^{{}^{\prime\prime}}_{\mathcal{S}j}\parallel}-\frac{A^{{}^{ \prime\prime}}_{\mathcal{T}j}}{\parallel A^{{}^{\prime\prime}}_{\mathcal{T}j} \parallel}\right\| \tag{5}

其中

A^{{}^{\prime\prime}}_{\mathcal{S}j}

和

A^{{}^{\prime\prime}}_{\mathcal{T}j}

分别代表学生网络和教师网络的第

个富含上下文信息的中间特征图。在计算差值矩阵之前，每个特征图沿着其通道进行归一化。

整体的损失函数是

L_{CE}

和

L_{Attn}

的加权总和，表示为：

L_{AttnFD}=L_{CE}+\alpha L_{Attn} \tag{6}

其中

\alpha

是一个权重系数，其调整方法在第4.2节中描述。在这种情况下，著名的交叉熵损失函数

L_{CE}

被用作学生网络预测和真实标签之间的分割损失。

4 Experimental Results

Datasets and Evaluation Metrics

已经在两个流行的语义分割数据集上进行了实验。Cityscapes 数据集专为理解城市场景而设计，包含2,975张训练图像、500张验证图像和1,525张测试图像，这些图像都经过了精心标注。它覆盖了30个类别，评估主要集中在19个类别上。PascalVOC 数据集由1,464张标记的训练图像、1,449张验证图像和1,456张测试图像组成，包括21个前景目标类别以及一个背景类别。

评估指标 根据标准，分割性能通过mIoU和像素准确度来评估，为了公平比较，这些指标是在三次运行中取平均的。模型的大小由报告的网络参数数量来指示。

Implementation Details

4.2.1 Network Architecture

为确保公平评估，实验采用了与文献[22]中描述相同的教师和学生网络。在所有实验中一致应用的教师网络是使用ResNet101作为 Backbone 网络的Deeplab V3+。在Deeplab V3+分割中的学生网络则采用了不同的 Backbone 网络，例如ResNet18和MobileNetv2。

4.2.2 Training Details

Pascal数据集和Cityscapes数据集都采用了类似的配置来训练学生网络。对于Pascal数据集，采用了批量大小为6，总共120个周期；而对于Cityscapes数据集，则采用了批量大小为4，总共50个周期。采用了随机梯度下降（SGD）优化器，Pascal的初始学习率设置为0.007，Cityscapes的初始学习率设置为0.01。学习率的调整基于余弦退火调度器。

在训练阶段之前，每张图像都要进行预处理，包括在原始尺寸的0.5到2倍之间进行随机缩放，水平随机翻转，以及随机裁剪到Pascal的

513\times 513

像素尺寸和Cityscapes的

512\times 1024

像素尺寸。教师网络和学生网络的主干部分都使用了来自ImageNet数据集的预训练权重，而分割部分则随机初始化。在方程6中描述的方法的唯一超参数是通过实验各种值来微调的。

对于Pascal数据集，它被设置为

\alpha=1

；对于Cityscapes数据集，它被设置为

\alpha=10

。在推理过程中，在单一尺度上对原始输入的性能进行评估。与使用原始特征图不同，按照[11]中的建议，使用了Pre-ReLU特征图，以保留应用注意力模块的负值。该实现是在PyTorch框架中进行的，所有网络都在单个NVIDIA GeForce RTX 3090 GPU上进行训练。

需要强调的是，通过引入卷积层来解决教师网络和学生网络在特征尺寸上的任何不一致。此外，在训练学生网络期间，教师网络的参数及其注意力模块先被训练，然后被冻结。

Results

进行了广泛的实验来评估所提出的AttnFD的性能。它与几种现有的蒸馏方法进行了比较；分别是：KD，AT，SP，ICKD，Overhaul，CIRKD，DistKD，以及LAD。上述每种方法都经过了所有主干网络、编码器和解码器特征的测试，以及最终输出图的测试，以确定最佳结果。

4.1.2 PascalVoc.

最初，将所提出方法的结果与上述方法在PascalVoc数据集上的结果进行比较。如表1所示，在没有蒸馏的情况下，AttnFD使模型的表现显著提升，当使用ResNet18作为学生 Backbone 网络时，性能提高了5.59%，使用MobileNet作为学生 Backbone 网络时，性能提升了6.46%。此外，作者的方法以较大优势超过了现有的最佳性能方法。

具体来说，与ResNet18作为学生 Backbone 网络时的LAD相比，作者的方法高出1.67%，并且与MobileNet作为学生 Backbone 网络时的CIRKD相比，提升了1.36%，使其成为第二佳方法。基于ResNet18的定性比较如图5所示。

4.1.3 Cityscapes.

所提出的方法与现有方法一起在Cityscapes数据集上进行了评估。表2中展示了定量结果。

表明AttnFD分别提升了ResNet18和MobileNet学生模型的性能8.95%和7.75%。值得注意的是，与DistKD和LAD这两个次优方法相比，AttnFD在mIoU上分别超过了它们1.23%和0.96%。此外，在像素准确度方面也有显著提升，比DistKD和LAD分别提高了2.08%和1.1%。利用ResNet18与DistKD进行的定性比较验证了所提方法的有效性，如图6所示。

Ablation Study

为了进一步证实所提出方法的有效性，进行了消融研究。由于该方法专注于最小化教师与学生细化特征图之间的MSE损失，作者探究了这些特征图在网络不同层的影响。表3展示了三种不同特征图的结果：“Backbone”代表通过网络 Backbone 获得低级特征，“Encoder”表示从编码器最后一层获得特征，“Decoder”指的是网络最终卷积层之前的特征图。

结果显示，所有这些特征都提高了性能，这在PascalVoc数据集上的ResNet18和MobileNet Backbone 上均得到了观察。值得注意的是，与来自Encoder的特征相比，来自Decoder的特征表现出更明显的改进，且两者都超过了来自Backbone的特征。这归因于网络最终层所包含的更丰富、更详细的信息。当这三种特征结合在一起时，它们共同达到了最佳性能。

图7展示了在PascalVoc数据集上，使用ResNet18作为学生模型时，AttnFD和LAD之间每个类别的mIoU对比。如图所示，AttnFD的表现大致与LAD持平或略好，在特定类别上有显著性能提升，例如自行车（+10.1）、羊（+5.96）、鸟（+4.85）和椅子（+4.1）。

同样地，图8展示了在Cityscapes数据集上，AttnFD与DistKD（同样使用ResNet18作为学生模型）之间的比较。AttnFD在火车（+12.91）和大巴（+2.82）等类别上展现出显著更优的性能。图6的顶行证实了这一点，突显了DistKD的。

将公交车误分类为火车的问题，通过AttnFD的性能提升得到了纠正。此外，与DistKD相比，AttnFD在交通灯和交通标志的分割上表现出色。这种改进可以归因于在细化特征图中增强了对微小目标（如交通灯）的突出显示，这有助于学生更有效地模仿教师的注意力，从而使得小目标的分割效果更好。

Discussion

提出的方法在分割中结合了额外的蒸馏损失以及交叉熵损失。与其它现有方法不同，这些方法执行多个蒸馏损失，如KD损失，它只需微调一个超参数。值得注意的是，尽管作者进行了各种尝试，将所提出的方法与KD损失结合，不仅没有在mIoU上得到任何提升，反而稍微降低了它（-0.5%）。这表明所提出的损失函数有效地从教师的精细化特征图中蒸馏出关键信息传递给学生。

此外，图7中显示的每个类别的mIoU结果证明，AttnFD在所有类别上始终超越或匹配KD方法的性能，表明结合KD损失并没有提供额外的益处。相比之下，LAD方法在特定类别（如自行车和飞机）上的表现不如KD方法，这突显了利用KD损失来提升这些特定类别在LAD方法中的结果的必要性。

此外，在训练后期减少作者的蒸馏损失系数的实验，如文献[28, 50]所证实的，这有助于提高性能，但作者的尝试却导致了性能的小幅下降。这表明CBAM模块有效地学习到了如何为传输到学生模型而强调关键信息，这一点从图1中得到支持，图1展示了噪声减少且重点区域突出的细化特征图，为蒸馏做好准备。

5 Conclusion

引入了一种新的语义分割方法。与现有方法不同，后者通常关注成对信息或涉及复杂的蒸馏损失，所提出的方法通过使用原始特征并应用卷积块注意力模块中的通道和空间注意力来简化过程，以细化特征图。这些细化特征突出了关键图像区域，并为蒸馏目的包含了丰富信息。在两个基准数据集上的大量实验一致表明，与没有蒸馏的模型相比，性能有显著提高。与最先进的方法的比较进一步验证了所提方法的有效性。