语义分割的归一化特征蒸馏

作者: Valar_Morghulis | 来源:发表于2022-08-28 18:09 被阅读0次

Normalized Feature Distillation for Semantic Segmentation

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

作者：Tao Liu, Xi Yang, and Chenshu Chen

单位：Hikvision Research Institute

日期：2022.7

摘要：作为一种很有前途的模型压缩方法，知识蒸馏通过从繁琐的模型中转移知识来提高紧凑模型的性能。用于指导学生训练的知识类型很重要。以前的语义分割蒸馏方法试图从特征中提取各种形式的知识，这涉及依赖于先验信息的精细手工设计，并且性能增益有限。在本文中，我们提出了一种简单而有效的特征蒸馏方法，称为归一化特征蒸馏（NFD），旨在实现对原始特征的有效蒸馏，而无需手动设计新的知识形式。关键思想是防止学生专注于模仿教师特征响应的幅度(magnitude)，我们通过归一化来实现这一点。我们的方法在城市景观、VOC 2012和ADE20K数据集上实现了最先进的语义分割蒸馏结果。代码将可用。

1 引言

最近关于主干网络[9,14,32,24]和分割框架[35,3,30]的工作大大提高了语义分割的性能。然而，这些高性能模型通常需要大量内存和计算开销。在实时应用中，由于资源有限，轻量级模型是首选的。设计轻量级网络[13,33,34,29]和执行模型压缩[19,7,12]是解决这个问题的可行方法。知识蒸馏（KD）是一种很有前途的模型压缩方法，它将繁琐模型（教师）的知识蒸馏为紧凑模型（学生）。

由于深度神经网络中间层的特征包含丰富的表征信息，特征蒸馏是知识发现中的一种常用方法。FitNets[20]将[12]的思想扩展到中间表示，以改进学生的训练。关键思想是让学生直接模仿教师的中间层特征。受此启发，提出了一系列方法，以间接方式对齐教师和学生之间的特征[6]。他们努力从特征中蒸馏各种形式的人工设计知识，如注意力图[31]、Gramian矩阵[28]和成对相似性[18]。这些方法通常依赖于传统特征工程中的先验知识，并且在对原始特征执行转换时遭受信息丢失，例如，特征的通道信息在[31]中丢失。尽管FitNets[20]的性能改进有限，但其优势在于没有信息丢失[11]。

为什么具有原始功能的蒸馏不能像预期的那样工作？我们认为，强迫学生模仿与教师完全相同的特征响应是低效的，在蒸馏中没有必要。因此，我们提出了一种简单而有效的语义分割特征蒸馏方法，称为归一化特征蒸馏（NFD），旨在通过蒸馏过程中的归一化来去除特征响应的大小(magnitude)。与以前的方法不同，我们期望学生以有效的方式从教师的原始特征中获取有价值的知识。因此，我们的方法不涉及阐述各种形式的知识。

我们的主要贡献总结如下：

• 我们提出了一种简单而有效的语义分割特征蒸馏方法，无需仔细设计各种形式的知识。

• 我们的经验表明，简单地鼓励学生从老师那里学习归一化特征分布可以产生优异的蒸馏性能。

• 我们在城市景观、VOC 2012和ADE20K数据集上评估了拟议方法。大量实验表明，我们的方法显著提高了基线学生模型的性能，实现了最先进的蒸馏结果。

2 相关工作

2.1 语义分割的知识蒸馏

知识蒸馏在图像分类领域得到了广泛的研究，已被证明是提高轻量级模型性能的一种很有前途的方法。随后的研究倾向于将知识蒸馏扩展到更具挑战性的密集预测任务，如语义分割[27,10,18,25,17,22]。直接将图像分类中的蒸馏方法应用于密集预测任务可能不会产生令人满意的结果。因此，提出了用于语义分割的蒸馏方法。

Xie等人[27]使用像素与其8个相邻像素之间的欧几里得距离来构建所谓的一阶知识，以确保学生和教师获得的分割边界信息彼此接近。Liu等人[18]通过计算特征图上的成对相似性来蒸馏长距离依赖性，并通过对抗性学习在教师和学生的输出之间实现高阶一致性。Wang等人[25]建议将教师的类别内特征变化迁移到学生身上。Shu等人[22]通过在教师和学生之间柔和地调整每个通道的激活来关注通道信息，这在logits上比在功能上更有效。

以前的大多数方法要么侧重于从特征中蒸馏各种形式的人工设计知识，要么主要侧重于logits蒸馏。与以前的方法不同，我们的目标是利用原始特征进行有效的蒸馏，而不是设计新的知识形式。

2.2 归一化

事实上，在训练神经网络的过程中，每个层的输入分布都会发生变化，这会减慢训练速度，并使深层神经网络的训练变得困难。Ioffe和Szegedy[15]将这一现象称为内部协变量移位，并提出批量归一化（BN），通过在小批量内计算的均值和方差对特征进行归一化来解决这一问题。它已被证明具有促进优化的能力，并能够使非常深的网络收敛。然而，“批量”的概念并不总是存在，或者可能会不时发生变化[26]。在此基础上，提出了一些不依赖批量维数的归一化方法。层归一化（LN）[1]沿通道维度操作，这对于递归神经网络非常有效。实例归一化（IN）[23]通常用于图像样式化，其工作方式与BN类似，但仅适用于每个样本。组标准化（GN）[26]将通道分成组，并在每组内计算标准化的平均值和方差，这在广泛的批量大小范围内始终有效。

在本文中，我们是通过沿着[15,1,23]（BN，LN，IN）的不同维度进行归一化来去除原始特征的幅度(magnitude)信息。但与它们[15,1,23]不同，我们方法中的归一化没有可训练的缩放和移位参数。

3 方法

3.1 初步

由于其简单有效，在[12]中提出的传统KD损失预测是许多视觉任务的常用目标。将语义分割视为每像素分类问题，以前的方法[18,25,22]采用这种损失，使学生模仿教师的每像素类概率，其可以表示为：

$\mathcal{L}_{kd}=D_\mathrm{KL}(\boldsymbol{Q}^t, \boldsymbol{Q}^s)$ （1）

其中 $\boldsymbol{Q}^t$ 和 $\boldsymbol{Q}^s$ 分别表示教师和学生生成的分割概率图， $D_\mathrm{KL}(\cdot)$ 是Kullback-Leibler散度。

让 $\boldsymbol{F}^t \in \mathbb{R}^{B \times C \times H \times W}$ 表示教师的特征，其中B是批量大小，C是通道数，H和W是高度和宽度。类似地， $\boldsymbol{F}^s \in \mathbb{R}^{B \times C^\prime \times H^\prime \times W^\prime}$ 表示学生的特征。特征蒸馏损失的一般形式可以表示为：

$\mathcal{L}_{feat}=D(T^t(\boldsymbol{F}^t), T^s(\boldsymbol{F}^s))$

其中 $T^t(\cdot)$ 和 $T^s(\cdot)$ 分别是应用于教师和学生特征的变换， $D(\cdot)$ 测量特征之间的距离，通常是L1或L2距离。由于通常的做法是通过卷积或上采样使 $\boldsymbol{F}^s$ 具有与 $\boldsymbol{F}^t$ 相同的尺寸，因此，如有必要，将在 $T^t(\cdot)$ 和 $T^s(\cdot)$ 之前进行尺寸对齐。不同的特征蒸馏方法通常采用不同的 $T^t(\cdot)$ 和 $T^s(\cdot)$ 。例如， $T^t(\cdot)$ 和 $T^s(\cdot)$ 可以是[31]中沿通道维度的特征求和。

3.2 归一化特征蒸馏

[20]中提出的朴素特征蒸馏鼓励学生直接模仿教师的原始特征。这意味着不需要等式（2）中的 $T^t(\cdot)$ 和 $T^s(\cdot)$ ，只需要学生特征的尺寸对齐。虽然这种方式没有遗漏任何信息，但其性能改进有限。我们以这种朴素的方式进行特征蒸馏，并分析教师和学生之间的特征差异。

首先，我们计算教师和学生每个通道的特征平均值，分别得到 $\mu^t \in \mathbb{R}^C$ 和 $\mu^s \in \mathbb{R}^C$ ，并计算以下误差

$e^{\mu}_i=\mu^t_i - \mu^s_i$

其中 $i=1,2,...,C$ 。 $e^{\mu}$ 的分布如图1（a）左侧所示。我们可以看到大多数通道的 $e^{\mu}$ 接近于零，这意味着原始特征蒸馏迫使学生和教师之间的特征响应幅度接近。

此外，我们采用CKA[16]，这是一种能够定量比较网络内和网络之间表示的相似性指数，用于测量教师和学生之间的特征相似性。CKA对表示的正交变换和各向同性缩放是不变的，因此特征响应的大小不影响CKA的相似性。我们分析每个通道的学生和教师特征之间的CKA相似性，如图1（b）的左部分所示。很明显，原始特征蒸馏不会导致学生和教师特征之间的高度相似性。

图1。在城市景观数据集上，原始特征蒸馏方法和我们方法的教师与学生的特征差异。图(a)中，我们分析每个通道的平均误差（公式(3)），图(b)中是教师和学生特征之间的CKA相似性。教师模型为PSPNet-R101，学生模型为PSPNet-R18。主干最后一层的特征用于蒸馏和计算CKA相似性。

上述实验表明，以朴素的方式蒸馏原始特征会导致学生专注于学习特征响应的幅度，而无法实现与教师的高特征相似性。基于这一观察，我们希望在蒸馏过程中忽略教师特征响应的幅度信息，并允许学生从教师的原始特征中有效学习。因此，我们提出了一种简单而有效的特征蒸馏方法，称为归一化特征蒸馏（NFD）。首先，我们对教师和学生的特征进行归一化，通过使其具有零均值和单位方差来去除其幅度信息。然后，我们最小化教师和学生的归一化特征之间的L2距离，类似于公式（2），其可以,公式化为：

$\mathcal{L}_{nfd}=D(Norm(\boldsymbol{F}_t), Norm(\boldsymbol{F}_s))$ （4）

等式（4）中的 $Norm$ 表示归一化：

$\boldsymbol{\hat{F}} = \frac{1}{\sigma}(\boldsymbol{F} - \mu)$ （5）

其中 $\boldsymbol{F}$ 是原始特征， $\boldsymbol{\hat{F}}$ 是归一化特征， $\mu$ 和 $\sigma$ 是特征的平均值和标准差。在归一化之前，我们将维度对齐应用于学生特征。受不同归一化方法的启发，我们可以沿着不同维度对特征进行归一化，例如BN[15]中的 $(B,H,W)$ ，LN[1]中的 $(C,H,W)$ 和IN[23]中的 $(H,W)$ 。与[15,1,23]不同，等式（5）中的归一化没有可训练的缩放和移位参数。我们将在第4.4节中比较和分析沿不同维度的归一化性能。默认情况下，我们沿 $(H,W)$ 维度对特征进行归一化。

与原始特征蒸馏相比，所提出的方法侧重于实现与教师的高相似度，而不是模仿特征响应的幅度，如图1（a）和图1（b）的右部分所示。

将我们的方法应用于语义分割的流程如图2所示。在前面的方法[18,25,22]的基础上，我们在Logit上也采用了传统的KD损失[12]。因此，我们方法的总损失可以表示为：

$\mathcal{L}= \mathcal{L}_{gt} + \lambda_1 \mathcal{L}_{kd} + \lambda_2 \mathcal{L}_{nfd}$ （6）

其中， $\mathcal{L}_{gt}$ 是语义分割的交叉熵损失， $\lambda_1$ 在[18,25,22]之后设置为10， $\lambda_2$ 设置为0.7。

图2。将我们的方法应用于语义分割的流程。

\mathcal{L}_{nfd}

是施加在中间层特征上的归一化特征蒸馏损失。

\mathcal{L}_{kd}

是Logit上的传统kd损失。

\mathcal{L}_{gt}

是语义分割的交叉熵损失

4 实验

4.1 数据集

城市景观 Cityscapes[4]是一个用于语义城市场景理解的大型数据集，除了一组较大的19998幅弱注释图像外，还具有5000幅图像的高质量像素级注释。它包含30个类，其中19个用于评估。精细注释的5000幅图像被分成2975、500和1525幅图像，用于训练、验证和测试。在实验中我们只使用精细注释的数据集。

PASCAL VOC 2012 PASCAL VOC 2012[5]数据集包含20个常见对象和一个背景类，并在日程拍摄的照片上添加注释。我们使用具有[8]提供的超粗注释的增强数据集，生成10582、1449和1456幅图像，用于训练、验证和测试。

ADE20K ADE20K[36]是一个注释密集的数据集，包含了stuff、objects和 parts的实例，涵盖了场景中的各种视觉概念。它包含150个类，分为20210、2000和3000个图像，用于训练、验证和测试。由于其类别数量众多，且在复杂场景中存在多个小对象，因此具有挑战性。

4.2 实施细节

网络架构。

我们采用带有ResNet101[9]主干的PSPNet[35]作为所有实验的教师模型。我们采用不同的分割模型（PSPNet[35]和DeepLabV3[2]）和主干（ResNet18[9]和MobileNetV2[21]）作为学生模型，以验证我们方法的有效性。

训练细节。

我们使用预训练的教师模型，并在蒸馏过程中保持其参数固定。对于学生的训练，我们使用随机梯度下降（SGD）作为优化器，批量大小为16，权重衰减为0.0005，动量为0.9。我们使用“poly”学习率策略，其中学习率等于 $base\_lr*(1 - \frac{iter}{max\_iter})^{power}$ 。我们将基本学习率设置为0.01，power设置为0.9。我们为Cityscapes和VOC 2012数据集训练80k迭代，为ADE20K数据集训练160k迭代。我们在训练期间对输入图像应用随机水平翻转、随机缩放（从0.5到2.0）和随机裁剪作为数据增强。城市景观、VOC 2012和ADE20K的裁剪大小分别为512×1024、512×512和512×512。我们对所有数据集使用单尺度测试。除非另有说明，在我们的方法中，主干最后一层的特征用于特征蒸馏。

4.3 与现有方法的比较

我们评估了提出的NFD，并将其与最近在城市景观[4]、PASCAL VOC 2012[5]和ADE20K[36]数据集语义分割中的知识蒸馏方法进行了比较。我们根据发布的代码重新实现了SKD[18]、IFVD[25]和CWD[22]。根据其推荐值设置与蒸馏损失相关的超参数。为了公平比较，包括NFD、SKD[18]、IFVD[25]和CWD[22]在内的所有方法都使用了与第4.2节所述完全相同的训练和测试策略。

表1显示了具有不同主干（ResNet18[9]和MobileNetV2[21]）和解码器（PPM[35]和ASPP[2]）的各种学生模型的结果。NFD显著提高了基线学生模型的性能。例如，在城市景观、VOC 2012和ADE20K中，NFD为PSPNet-R18带来的性能增益分别为3.98%、4.03%和4.06%。虽然NFD默认情况下利用主干网络最后一层的功能进行蒸馏，但性能增益受主干架构的影响不大。具体而言，NFD在城市景观、VOC 2012和ADE20K上分别将PSPNet-MV2的性能提高了2.77%、4.29%和3.68%。此外，NFD进一步缩小了教师模型与DeepLabV3-R18之间的差距，DeepLabV4-R18作为一个强大的基线学生模型。

更重要的是，在各种实验设置下，特别是在VOC 2012和ADE20K数据集上，所提出的NFD始终优于其他方法。例如，在更具挑战性的ADE20K数据集上，当使用PSPNet-R18、PSPNet-MV2、DeepLabV3-R18和DeepLabV3-MV2作为学生时，NFD比CWD[22]要好2.31%、1.75%、1.73%和2.45%。

4.4 消融研究

在本节中，我们进行了大量实验，以验证我们方法的有效性，并讨论了一些超参数的选择。消融实验主要在城市景观和ADE20K数据集上进行，PSPNet-R101作为教师模型，PSPNet-R18作为学生模型。

特征归一化。

如前所述，我们可以在执行NFD时沿不同维度对特征进行归一化，例如BN[15]中的 $(B,H,W)$ ，LN[1]中的 $(C,H,W)$ ，IN[23]中的 $(H,W)$ 。我们进行实验来研究不同归一化维数的影响。如表2所示，所提出的NFD在不同归一化维度下实现了类似的性能，并且所有归一化设置都优于没有归一化的情况。无归一化的NFD相当于[20]中的原始特征蒸馏。

为了进一步研究NFD对特征蒸馏的影响，我们使用CKA[16]测量主干网络最后一层教师和学生特征之间的相似性。我们从城市景观验证集中随机抽取100幅图像，计算CKA相似度。在蒸馏过程中，学生特征有维度对齐，因此该层中的教师和学生特征都有2048个通道。我们计算每个通道的相似度，获得2048个CKA相似度。

为了进行定性分析，我们将2048个通道的CKA相似性reshape为32×64的矩阵，以便以热图的形式进行更好的可视化，如图3所示。它清楚地表明，无论沿着哪个维度进行归一化，所提出的NFD总是导致学生和教师特征之间的更高相似性。

对于定量分析，我们分析了2048个通道的CKA相似性分布，如图4所示。对于建议的NFD，大多数通道的CKB相似性在[0.8，1.0]的范围内。相反，对于原始特征蒸馏，CKA相似性主要分布在[0.6,0.8]的范围内。此外，我们在表3中计算了所有通道的平均CKA相似度。结果表明，建议的NFD使学生能够更好地模仿老师，从而获得更高的特征相似度。

特征蒸馏的位置。

在特征蒸馏期间，特征可以来自学生和教师网络的任何可用层。这里我们选择了三个常用的蒸馏位置：1）主干的最后一层，2）解码器的最后一个层（例如，PPM[35]），以及3）最终预测层。在之前的实验中，我们默认使用主干最后一层的特征，现在我们在表4中展示了不同蒸馏位置的NFD结果。我们的方法在不同位置，特别是主干最后一个层，效果良好。解码器或预测层产生的特征是高度压缩的，并且是特定于任务的，因为它们位于模型的顶层。相反，主干中的特征往往更具信息性，并且具有丰富的通用表示，这可以解释主干中NFD的更好结果。这一结果表明，我们的方法在其他视觉任务中具有很大的知识蒸馏潜力。

权重。

我们方法中的特征蒸馏损失由等式（6）中的 $\lambda_2$ 加权。在城市景观、VOC 2012和ADE20K数据集上进行了广泛的实验，以研究拟议的NFD对 $\lambda_2$ 的敏感性。我们使用 $\lambda_2$ ∈ {0.3, 0.5, 0.7, 1.0, 1.3, 1.5}，每个实验运行3次。图5中的结果证明了所提出的NFD对超参数的优异鲁棒性。

不同形式的知识。

先前的方法设计了从特征中蒸馏的各种形式的知识，如成对相似性[18]、类内特征变化[25]和通道概率[22]。本质上，我们的方法并不是设计一种新的知识形式，而是从原始特征中去除幅度信息，我们认为这对于蒸馏是不必要的。为了公平地比较不同形式的知识对特征蒸馏的影响，我们从所有方法中删除与特征蒸馏无关的部分来进行实验。具体而言，我们删除了[18,25,22]中的传统KD损失[12]和建议的NFD，并删除了[19,25,23]中强加给Logit的对抗性蒸馏损失。结果如表5所示，其中所有方法都使用主干网络最后一层的特征进行蒸馏。在这个实验装置中，我们的方法的优越性更加明显，证明了它的简单性和有效性。

5 结论

在本文中，我们提出了一种简单而有效的语义分割特征蒸馏方法，称为归一化特征蒸馏（NFD）。与以前的方法不同，这些方法力求手工设计各种形式的知识，我们关注的是如何在不需要设计新形式的知识的情况下使具有原始特征的蒸馏有效。我们认为，强迫学生模仿与教师完全相同的特征反应是低效的，在蒸馏中不必要，反而会增加学习难度。因此，提出的NFD旨在通过蒸馏过程中的归一化来从原始特征中去除幅度信息。大量的语义分割实验表明，我们的NFD可以显著提高基线学生模型的性能。在未来的工作中，我们将把所提出的方法应用于其他视觉任务，如图像分类、目标检测和实例分割。此外，我们考虑将我们的方法应用于Transformer网络。

语义分割的归一化特征蒸馏

相关文章

网友评论

延伸阅读

深度阅读

栏目导航

热点阅读