在青源LIVE第31期中,旷视研究院张祥雨团队的张培圳研究员深入浅出地为我们介绍了其团队被 NeurIPS 2021 录用的论文「用于物体检测的实例条件知识蒸馏」。本期报告首先简要回顾了用于图像识别、物体检测的知识蒸馏技术的发展历程,介绍了论文的研究动机、模型架构、训练技巧以及实验结果。
本文整理自青源 LIVE 第 31 期,视频回放链接:
https://event.baai.ac.cn/activities/204
01 背景知识
早在 NIPS 2014 上,Hinton 等人在论文「Distilling the knowledge in a Neural Network」中就开始研究用于图像分类的知识蒸馏技术。根据知识源的不同,我们可以将用于图像分类的知识蒸馏技术分为「基于 Logits 响应」、「基于中间特征图激活」和「基于关系建模」的知识蒸馏。
如上图所示,基于响应的知识蒸馏往往会在弱监督场景下计算预训练教师网络和学生网络用于分类的 Logits 的 KL 散度;基于激活的知识蒸馏会计算教师网络和学生网络中间特征图的 MSE;基于关系建模的知识蒸馏则会针对教师网络和学生网络对每一个 mini-batch 中图片经过全局平均池化后的特征矩阵计算 MSE 等距离。
面向物体检测任务的知识蒸馏技术大多属于基于中间特征图激活的知识蒸馏。近年来,一些面向工程化的工作也会尝试系统地使用上述知识蒸馏方式的组合(例如,FBKD、GID、DeFeat、G-DetKD,等)
相较之下,分类模型需要处理的图像往往包含单个物体,待分类的物体位于图像的中央;而检测模型需要处理的图像中往往包含多个物体,这些物体从属于各种类别、分布于图像的各个位置,图像中还包含大量冗余的背景信息。
在针对中间特征图进行蒸馏时,蒸馏特征的位置是十分关键的,我们通常会利用以下三类拓扑信息:
(1)前景先验:给定训练图片中物体的边界框 BBOX,我们将其映射到中间特征图上,对该区域进行知识蒸馏;对于基于 anchor 的物体检测任务,重点对与 groundtruth 重叠比例较高的 anchor 区域进行蒸馏;对于多阶段物体检测任务,针对 proposal 进行蒸馏;
(2)前景和背景先验:将前景、背景分开蒸馏/在全图上使用注意力机制;
(3)基于度量的区域:根据设计出的度量方式找出有效的蒸馏信息(例如,教师网络与学生网络的差异)。
如上图所示,蓝色和棕色的实线方框代表真实的物体边界框;黑色的方框代表物体附近的 anchor;蓝色、棕色、红色的虚线方框代表多阶段检测模型提出的 proposals,可以通过非极大抑制等手段对其进行修正。
在考虑背景先验时,我们可以通过二值掩膜将前景和背景分隔开,进而分别进行蒸馏。此外,我们可以在全图中分别在空间和通道上使用注意力机制,我们可以分别对教师和学生网络计算空间维度和通道维度上的注意力图,进而对这些注意力两两进行知识蒸馏。
在利用教师网络、学生网络评分差异进行知识蒸馏时,假设教师网络和学生网络预测头的预测特征图的维度相同,图像上的某个像素点邻近于某个 anchor 的预测框,模型会计算出该预测框从属于某一类的置信度以及物体回归的位置编码信息。在典型的情况下,教师网络可以很好地预测出类别和边框位置信息,而学生网络对于某些类别物体的预测性能较差。因此,我们可以基于度量学习的方式,重点对学生网络学习性能较差的类别进行蒸馏。
以 ADL 模型为例,其关键之处在于计算自适应蒸馏权重(ADW)。如上图所示,红色圆圈代表教师网络和学生网络在某个空间位置上预测评分的 KL 散度。此外,在 GID 模型中,作者定义了一种通用实例(General Instance)代表教师网络或学生网络的 proposal,这两种网络对proposal 会产生预测得分的差异,GI Score 为针对所有类的得分差异的最大值,该值越大说明不确定性越强。
基于前景/背景的建模方法非常直接,符合人类的认知习惯,但是这些方法获得的知识可能是次优的。基于度量的方法相对来说较为灵活,但是往往并不直观,并且存在多种可能的度量方法。
02 用于物体检测的实例条件知识蒸馏
为了解决上述问题,旷视研究院张祥雨组提出了「用于物体检测的实例条件知识蒸馏」,在保证模型具有直观性、可解释性的同时,兼具了基于度量的蒸馏方法的灵活性。模型的整体架构如上图所示。
具体而言,该方法利用人类标注的实例作为条件,来查询教师网络的特征,从而检索出有用的知识,用于蒸馏学生网络的特征图。
如上图所示,左侧为 Query 的特征构造过程,右侧为抽象的建模过程。从抽象建模过程开始,我们可以将图中物体的 one-hot 类别信息和实例的位置嵌入通过多层感知机编码到 Query 中。这样一来,每个实例的特征都可以被表示为一个低维向量。
接着,我们通过「键-值」形式的注意力,将物体标注和保留了空间-拓扑信息的特征图像素之间的关系实例化,从而得到了「实例感知注意力掩模」。这样一来,我们就可以衡量物体和图像位置的关系。其中,基于 Query、Key、Value 的自注意力的计算范式与原始的 Transformer 架构中的计算范式相同。
为了保证掩模与待检测物体的位置相关,我们设计了面向检测任务的辅助任务引导对掩模的学习。具体而言,我们利用注意力机制的输出作为识别、定位辅助任务的监督信号。值得注意的是,辅助任务不会更新学生网络的参数。在识别任务中,我们首先根据数据集先验采样出一些伪样本。对于识别出的真实样本,我们进一步对其位置进行定位。
03 实验结果
与现有的知识蒸馏方法相比,本文提出的方法在 COCO 数据集上取得了较优的性能。在使用 RetinaNet 作为物体检测网络,采用多尺度训练时,本文提出的方法可以将学生网络(ResNet-50)的物体检测 AP 从 37.4 提升到 40.7,这一数字甚至超过了教师网络(ResNet-101)的 40.4 的 AP。在训练过程中进行的,涉及物体Groundtruth信息的实例建模为学生从单纯的教师网络特征图信息以外获得更好训练,提供了可能。
在其它的检测和分割模型上,本文提出的知识蒸馏方法也取得了优秀的性能表现。如上图所示,继承策略(Inheriting Strategy)是一种非常有用的训练技巧:当教师网络和学生网络的Neck部分和检测头部分的架构相同时,学生网络就可以继承上面两个部分的参数进行良好的初始化,从而进一步提升性能或快速收敛。
在使用 EfficientNet-B0、MobileNet-V2 等面向移动设备的主干网络时,本文提出的知识蒸馏方法可以大大提升物体检测的 AP。
Transformer 解码器的注意力头在原图中的激活可视化结果如上图所示。实验结果表明,注意力头可以有效注意到实例周围的信息。
04 讨论
在整体的计算过程中,计算量主要体现在教师网络的预训练和学生网络本身的训练开销,而本文加入的辅助任务和知识蒸馏的计算开销相对较小。在知识迁移方面,本文提出的方法可以直观地检索图像中有意义的内容,引入了类似基于度量的蒸馏方法的灵活特性,具有一定的可解释性。本文采用的继承策略训练技巧可以有效地加快收敛。
05 结语
在本文中,旷视研究院提出了一种面向物体检测的新型知识蒸馏框架,该框架可以学习特征解码器模块从而检索有用的知识,设计了辅助任务来优化解码器。作者设计了基于实例的条件,从而适应目标检测任务,其中 Query 为标注的实例信息,Key 和 Value 为教师网络的表征知识。
首发:智源社区
作者:张培圳
专栏文章推荐
- 旷视开源 MegFile——Python 文件读写 SDK
- 开源深度学习框架项目参与指北 - 内含易上手任务清单
- 旷视研究院夺得 NeurIPS 2021 ML4CO 组合优化比赛 Dual task 赛道第一
欢迎关注旷视研究院极术社区专栏,定期更新最新旷视研究院成果
加入旷视:career@megvii.com