在计算机视觉方面,实现最先进性能的大型模型与实际应用中简单的模型之间的差距越来越大。在本文中,将解决这个问题,并显著地弥补这2种模型之间的差距。
在实证研究中,作者的目标不是一定要提出一种新的方法,而是努力确定一种稳健和有效的配置方案,使最先进的大模型在实践中能够得到应用。本文证明了在正确使用的情况下,知识蒸馏可以在不影响大模型性能的情况下减小它们的规模。作者还发现有某些隐式的设计选择可能会极大地影响蒸馏的有效性。
作者的主要贡献是明确地识别了这些设计选择。作者通过一项全面的实证研究来支持本文的发现,在广泛的视觉数据集上展示了很不错的结果,特别是,为ImageNet获得了最先进的
ResNet-50
模型,达到了82.8%的Top-1精度。
1简介
大型视觉模型目前主导着计算机视觉的许多领域。最新的图像分类、目标检测或语义分割模型都将模型的大小推到现代硬件允许的极限。尽管它们的性能令人印象深刻,但由于计算成本高,这些模型很少在实践中使用。
相反,实践者通常使用更小的模型,如ResNet-50
或MobileNet
等,这些模型运行起来代价更低。根据Tensorflow Hub
的5个BiT
的下载次数,最小的ResNet-50
的下载次数明显多于较大的模型。因此,许多最近在视觉方面的改进并没有转化为现实世界的应用程序。
为了解决这个问题,本文将专注于以下任务:给定一个特定的应用程序和一个在它上性能很好的大模型,目标是在不影响性能的情况下将模型压缩到一个更小、更高效的模型体系结构。针对这个任务有2种广泛使用的范例:模型剪枝
和知识蒸馏
。
模型剪枝
通过剥离大模型的各个部分来减少大模型的大小。这个过程在实践中可能会有限制性:首先,它不允许更改模型族,比如从ResNet
到MobileNet
。其次,可能存在依赖于架构的挑战;例如,如果大模型使用GN
,修剪通道可能导致需要动态地重新分配通道组。
相反,作者专注于没有这些缺点的知识蒸馏
方法。知识蒸馏
背后的理念是“提炼”一个教师模型
,在本文例子中,一个庞大而繁琐的模型或模型集合,制成一个小而高效的学生模型
。这是通过强迫学生模型的预测与教师模型的预测相匹配,从而自然地允许模型家族的变化作为压缩的一部分。
图1
密切遵循Hinton
的原始蒸馏配置,发现如果操作正确,它惊人地有效;如图1所示作者将蒸馏解释为匹配教师和学生实现的函数的任务。通过这种解释发现对模型压缩的知识蒸馏的2个关键原则。
- 首先,教师和学生模型应该处理完全相同的输入图像,或者更具体地说,相同的裁剪和数据增强;
- 其次,希望函数在大量的支撑点上匹配,以便更好地推广。
使用Mixup
的变体,可以在原始图像流形外生成支撑点。考虑到这一点,通过实验证明,一致的图像视图、合适的数据增强和非常长的训练计划是通过知识蒸馏
使模型压缩在实践中工作良好的关键
。
尽管发现明显很简单,但有很多种原因可能会阻止研究人员(和从业者)做出建议的设计选择。
- 首先,很容易预先计算教师对离线图像的激活量,以节省计算量,特别是对于非常大的教师模型;
- 其次,
知识蒸馏
也通常用于不同的上下文(除了模型压缩),其中作者推荐不同甚至相反的设计选择; - 组后,
知识蒸馏
需要比较多的Epoch来达到最佳性能,比通常用于监督训练的Epoch要多得多。更糟糕的是,在常规时间的训练中看起来不理想的选择往往是最好的,反之亦然。
在本文的实证研究中,主要集中于压缩大型BiT-ResNet-152x2
,它在ImageNet-21k
数据集上预训练,并对感兴趣的相关数据集进行微调。在不影响精度的情况下,将其蒸馏为标准的ResNet-50
架构(用GN
代替BN
)。还在ImageNet
数据集上取得了非常强的结果:总共有9600个蒸馏周期,在ImageNet
上得到了新的ResNet-50
SOTA结果,达到了惊人的82.8%。这比原始的ResNet-50
模型高出4.4%,比文献中最好的ResNet-50
模型高出2.2%。
最后,作者还证明了本文的蒸馏方案在同时压缩和更改模型时也可以工作,例如BiT-ResNet
架构到MobileNet
架构。
2实验配置
2.1 Datasets, metrics and evaluation protocol
在5个流行的图像分类数据集上进行了实验:flowers102
,pets
,food101
,sun397
和ILSVRC-2012
(“ImageNet”)。这些数据集跨越了不同的图像分类场景;特别是,它们的类的数量不同,从37到1000个类,训练图像的总数从1020到1281167个不等。
2.2 Teacher and student models
在本文中,选择使用来自BiT
的预训练教师模型,该模型提供了大量在ILSVRC-2012
和ImageNet-21k
数据集上预训练的ResNet
模型,具有最先进的准确性。BiT-ResNets
与标准ResNets
唯一显著的区别是使用了GN层
和权重标准化
。
特别地专注于BiT-M-R152x2
架构:在ImageNet-21k
上预训练的BiT-ResNet-152x2
(152层,“x2
”表示宽度倍数)。该模型在各种视觉基准上都显示出了优异的性能,而且它仍然可以使用它进行广泛的消融研究。尽管如此,它的部署成本还是很昂贵的(它需要比标准ResNet-50
多10倍的计算量),因此该模型的有效压缩具有实际的重要性。对于学生模型的架构,使用了一个BiT-ResNet-50
变体,为了简洁起见,它被称为ResNet-50
。
2.3 Distillation loss
这里使用教师模型的和学生模型的之间的KL散度
作为一个蒸馏损失来预测类概率向量。对于原始数据集的硬标签,不使用任何额外的损失:
C是类别。这里还引入了一个温度参数T,用于在损失计算之前调整预测的softmax-probability
分布的熵:
2.4 Training setup
为了优化,使用带有默认参数的Adam优化器
训练模型。还使用了不带有Warm up
的余弦学习率机制。
作者同时还为所有的实验使用了解耦的权重衰减机制。为了稳定训练,在梯度的全局l2范数上以1.0的阈值进行梯度裁剪。最后,除在ImageNet
上训练的模型使用batch size
为4096进行训练外,对其他所有实验都使用batch size
为512。
本文的方案的另一个重要组成部分是Mixup
数据增强策略。特别在“函数匹配”策略中中引入了一个Mixup
变量,其中使用从[0,1]均匀抽样的较强的Mixup
系数,这可以看作是最初提出的β分布抽样的一个极端情况。
作者还使用了““inception-style”的裁剪,然后将图像的大小调整为固定的正方形大小。此外,为了能够广泛的分析在计算上的可行(训练了数十万个模型),除了ImageNet
实验,使用标准输入224×224分辨率,其他数据集均使用相对较低的输入分辨率,并将输入图像的大小调整为128×128大小。