CVPR 2022 | 用于语义分割任务的跨图像关系型知识蒸馏

KD作为一种思想已被广泛应用，尤其是图像分类任务。但是先前的工作^[1]表明直接将图像分类KD方法直接迁移到语义分割任务是不可行的，因为它们没有考虑到语义分割的结构化信息。一些经典的语义分割KD工作于是去尝试捕捉到一些上下文信息来建模结构化知识，比如pixel相似度^[1]，pixel与region向量的相似度关系^[2]，但是这些方法通常在单张图像内部进行信息提取，难以捕捉到更加广泛的上下文依赖。

于是，我们提出一种跨图像的知识蒸馏方法来建模图像间的依赖，从而使得网络能够捕捉到更加丰富的结构化信息。具体地，本工作建模数据集中的全部图像之间建模pixel-to-pixel以及pixel-to-region的对比关系作为知识。动机在于一个好的教师模型可以产生更好的全局pixel依赖以及特征空间。CIRKD引导学生模型去模仿教师网络产生的更好的结构化语义关系，因此提升了语义分割表现。

本方法分别从mini-batch和memory的角度对结构化关系进行建模。

1.  基于mini-batch的pixel-to-pixel蒸馏

首先是在batch中建模关系，将图像产生的特征图在spatial维度上进行分离，产生pixel-wise的特征向量。针对某一个pixel点的特征向量作为锚样本，其他pixel点特征向量（包括当前图像和其他图像）作为对比样本，可以得到相似度矩阵。基于此方法，在教师端和学生端都可以产生对应的相似度矩阵，然后通过KL散度进行逼近，使得学生可以学到教师的相似度矩阵：

一个batch内的图像进行两两之间的关系矩阵求取和迁移：

整个过程的示意图如下所示：

2.  基于memory的pixel-to-pixel蒸馏

相对于基于mini-batch的对比性关系求取，这里采用一个memory bank来存储对比向量，使得每一个锚样本都可以得到充足的对比样本。针对某一个pixel点的特征向量作为锚样本，每次训练随机从memory bank中采样大量的pixel特征向量来构造对比样本，通过相乘分别得到教师和学生的相似度矩阵：

使用KL散度的误差形式将教师和学生的pixel-to-pixel相似度矩阵进行对齐：

3.  基于memory的pixel-to-region蒸馏

离散化的pixel特征向量不能充分在捕捉到图像内容，因此本方法利用memory对region向量进行存储。region向量产生自对来自相同类别的pixel向量做一个平均。针对某一个pixel点的特征向量作为锚样本，每次训练随机从memory bank中采样大量的region特征向量来构造对比样本，通过相乘分别得到教师和学生的相似度矩阵：

使用KL散度的误差形式将教师和学生的pixel-to-region相似度矩阵进行对齐：

基于memory进行KD的整体示意图如下所示：

最终的训练误差则是将以上3部分误差相加在一起：

实验在流行的Cityscapes场景分割数据集上进行，相比于其他流行的语义分割KD方法，CIRKD获得了一致最佳的表现。

在一些case study上进行可视化，本方法能够得到最好的分割表现。

参考文献

[1] Yifan Liu, Ke Chen, Chris Liu, Zengchang Qin, Zhenbo Luo, and Jingdong Wang. Structured knowledge distillation for semantic segmentation. CVPR, 2019.

[2] Yukang Wang, Wei Zhou, Tao Jiang, Xiang Bai, and Yongchao Xu. Intra-class feature variation distillation for semantic segmentation. ECCV, 2020.