python | 实现带Temperature的SoftMax(knowledge distillation)

Python013

python | 实现带Temperature的SoftMax(knowledge distillation),第1张

在Hiton大佬的知识蒸馏文章中提出用带T(Temperature)的softmax,来达到在各个类上的概率分布更加soft的效果。

这里主要是简单实现了一下。

比较简单,就当记录一下吧~~

“很多昆虫在幼虫形态的时候是最擅长从环境中吸取能量和养分的,而当他们成长为成虫的时候则需要擅长完全不同能力比如迁移和繁殖。”在2014年Hinton发表的知识蒸馏的论文中用了这样一个很形象的比喻来说明知识蒸馏的目的。在大型的机器学习任务中,我们也用两个不同的阶段 training stage 和 deployment stage 来表达两种不同的需求。training stage(训练阶段)可以利用大量的计算资源不需要实时响应,利用大量的数据进行训练。但是在deployment stage (部署阶段)则会有很多限制,比如计算资源,计算速度要求等。知识蒸馏就是为了满足这种需求而设计的一种模型压缩的方法。

知识蒸馏的概念最早是在2006年由Bulica提出的,在2014年Hinton对知识蒸馏做了归纳和发展。知识蒸馏的主要思想是训练一个小的网络模型来模仿一个预先训练好的大型网络或者集成的网络。这种训练模式又被称为 "teacher-student",大型的网络是“老师”,小型的网络是“学生”。

在知识蒸馏中,老师将知识传授给学生的方法是:在训练学生的过程中最小化一个以老师预测结果的概率分布为目标的损失函数。老师预测的概率分布就是老师模型的最后的softmax函数层的输出,然而,在很多情况下传统的softmax层的输出,正确的分类的概率值非常大,而其他分类的概率值几乎接近于0。因此,这样并不会比原始的数据集提供更多有用的信息,没有利用到老师强大的泛化性能,比如,训练MNIST任务中数字‘3’相对于数字‘5’与数字‘8’的关系更加紧密。为了解决这个问题,Hinton在2015年发表的论文中提出了‘softmax temperature’的概念,对softmax函数做了改进:

这里的就是指 temperature 参数。当等于1 时就是标准的softmax函数。当增大时,softmax输出的概率分布就会变得更加 soft(平滑),这样就可以利用到老师模型的更多信息(老师觉得哪些类别更接近于要预测的类别)。Hinton将这样的蕴含在老师模型中的信息称之为 "dark knowledge",蒸馏的方法就是要将这些 "dark knowledge" 传给学生模型。在训练学生的时候,学生的softmax函数使用与老师的相同的,损失函数以老师输出的软标签为目标。这样的损失函数我们称为"distillation loss"。

在Hinton的论文中,还发现了在训练过程加上正确的数据标签(hard label)会使效果更好。具体方法是,在计算distillation loss的同时,我利用hard label 把标准的损失( )也计算出来,这个损失我们称之为 "student loss"。将两种 loss 整合的公式如下:

这里的是输入, 是学生模型的参数, 是交叉熵损失函数, 是 hard label , 是参数有的函数, 是系数, 分别是学生和老师的logits输出。模型的具体结构如下图所示:

在上述公式中, 是作为超参数人为设置的,Hinton的论文中使用的的范围为1到20,他们通过实验发现,当学生模型相对于老师模型非常小的时候, 的值相对小一点效果更好。这样的结果直观的理解就是,如果增加的值,软标签的分布蕴含的信息越多导致一个小的模型无法"捕捉"所有信息但是这也只是一种假设,还没有明确的方法来衡量一个网络“捕捉”信息的能力。关于,Hinton的论文中对两个loss用了加权平均: 。他们实验发现,在普通情况下 相对于非常小的情况下能得到最好的效果。其他人也做了一些实验没用加权平均,将设置为1,而对进行调整。

Hinton的论文中做了三个实验,前两个是MNIST和语音识别,在这两个实验中通过知识蒸馏得到的学生模型都达到了与老师模型相近的效果,相对于直接在原始数据集上训练的相同的模型在准确率上都有很大的提高。下面主要讲述第三个比较创新的实验:将知识蒸馏应用在训练集成模型中。

训练集成模型(训练多个同样的模型然后集成得到更好的泛化效果)是利用并行计算的非常简单的方法,但是当数据集很大种类很多的时候就会产生巨大的计算量而且效果也不好。Hinton在论文中利用soft label的技巧设计了一种集成模型降低了计算量又取得了很好的效果。这个模型包含两种小模型:generalist model 和 specialist model(网络模型相同,分工不同)整个模型由很多个specialist model 和一个generalist model 集成。顾名思义generalist model 是负责将数据进行粗略的区分(将相似的图片归为一类),而specialist model(专家模型)则负责将相似的图片进行更细致的分类。这样的操作也非常符合人类的大脑的思维方式先进行大类的区分再进行具体分类,下面我们看这个实验的具体细节。

实验所用的数据集是谷歌内部的JFT数据集,JFT数据集非常大,有一亿张图片和15000个类别。实验中 generalist model 是用所有数据集进行训练的,有15000个输出,也就是每个类别都有一个输出概率。将数据集进行分类则是用Online k-means聚类的方法对每张图片输入generalist model后得到的软标签进行聚类,最终将3%的数据为一组分发给各个specialist,每个小数据集包含一些聚集的图片,也就是generalist认为相近的图片。

在specialist model的训练阶段,模型的参数在初始化的时候是完全复制的generalist中的数值(specialist和generalist的结构是一模一样的),这样可以保留generalist模型的所有知识,然后specialist对分配的数据集进行hard label训练。但是问题是,specialist如果只专注于分配的数据集(只对分配的数据集训练)整个网络很快就会过拟合于分配的数据集上,所以Hinton提出的方法是用一半的时间进行hard label训练,另一半的时间用知识蒸馏的方法学习generalist生成的soft label。这样specialist就是花一半的时间在进行小分类的学习,另一半的时间是在模仿generalist的行为。

整个模型的预测也与往常不同。在做top-1分类的时候分为以下两步:

第一步:将图片输入generalist model 得到输出的概率分布,取概率最大的类别k。

第二步:取出数据集包含类别k的所有specialists,为集合 (各个数据集之间是有类别重合的)。然后求解能使如下公式最小化的概率分布q作为预测分布。

这里的KL是指KL散度(用于刻画两个概率分布之间的差距) 和 分别是测试图片输入generalist 和specialists(m)之后输出的概率分布,累加就是考虑所有属于 集合的specialist的“意见”。

由于Specialist model的训练数据集很小,所以需要训练的时间很短,从传统方法需要的几周时间减少到几天。下图是在训练好generalist模型之后逐个增加specialist进行训练的测试结果:

从图中可以看出,specialist个数的增加使top1准确个数有明显的提高。

本文结合Hinton在2014年发表的论文对知识蒸馏和相关实验做了一个简单的介绍,如今很多模型都用到了知识蒸馏的方法,但知识蒸馏在深度学习中还是非常新的方向,还有非常多的应用场景等待研究。

项目地址: https://momodel.cn/explore/5dc3b1223752d662e35925a3?type=app

[1]Hinton G, Vinyals O, Dean J. Distilling the knowledge in a neural network[J]. arXiv preprint arXiv:1503.02531, 2015.

[2] https://nervanasystems.github.io/distiller/knowledge_distillation.html

[3] https://www.youtube.com/watch?v=EK61htlw8hY&t=3323s

最近利用知识蒸馏的方法,对业务中的性能有了可观的提升,因此在这里总结一波。本文主要从宏观的角度分析一下各个蒸馏算法的蒸馏方式,具体细节可以根据兴趣阅读论文~ 知识蒸馏是一种模型压缩常见方法,用于模型压缩指的是在teacher-student框架中,将复杂、学习能力强的网络学到的特征表示“知识蒸馏”出来,传递给参数量小、学习能力弱的网络。从而我们会得到一个速度快,能力强的网络,因此这是一个概念上的模型压缩方案。从另一个角度来说,蒸馏可以使得student学习到teacher中更加软化的知识,这里面包含了类别间的信息,这是传统one-hot label中所没有的。由于蒸馏中软化标签的本质,因此蒸馏也可以被认为是一种正则化的策略。总结来说,知识蒸馏除了能够学习到大模型的特征表征能力,也能学习到one-hot label中不存在的类别间信息。现有的知识蒸馏方法主要侧重于两点: 从teacher的什么位置学习 和 用什么方式学习 。以下的总结图概述了本文要介绍的蒸馏方法。

目录结构:

(1)KL:知识蒸馏:蒸馏开山之作

https://arxiv.org/pdf/1503.02531.pdf

如上图所示,本文中直接利用KL散度来衡量教师模型和学生模型的输出分布,通过最小化KL散度的方式,使得学生模型的输出分布能够尽可能的逼近教师模型,从而实现知识蒸馏的目的。KL散度是一种衡量两个概率分布之间的差异的数学概念,有不懂的同学请出门左拐百度一下,右拐也行Google一下。

(2)FT:相关性因子加权学习法

https://arxiv.org/pdf/1802.04977.pdf

(3)PKT:概率分布学习法

https://arxiv.org/pdf/1803.10837.pdf

上述两篇文章的作者认为学生一般都是不聪明的,为了让学生能够更好的理解教师模型,FT算法这篇文章提出了一种新的知识转移方式,如图所示,利用卷积运算对教师模型的输出进行编码,并解码(翻译)给学生。而位于学生模块部分也添加一个卷积操作,用来学习翻译后的教师知识。实验证明这种方式要比直接学习效果好。PKT算法这篇文章提出了另一种新的知识转移方式,如图所示,该文章让学生模型学习教师模型的概率分布,使得整体的学习更加容易,更鲁棒。作者提出了一种通过匹配数据在特征空间中的概率分布进行知识蒸馏,PKT算法的另一个优势是该方法可以直接转移不同架构和维度层之间的知识。

(4)RKD:关系型学习法

https://arxiv.org/pdf/1904.05068.pdf

(5)CC:多输入联系型学习法

https://arxiv.org/pdf/1904.01802.pdf

所谓的单打独斗就是一个样本进行自我学习,单打独斗的蒸馏方法使得学生模型只能学习教师模型的输出表现,无法真正学习到教师模型的结构信息。而这两篇文章的作者都提出了多个样本之间进行合作学习的蒸馏学习方法,使得学生模型能够更好的学习到教师模型的结构信息。RKD关系型学习算法的核心是以多个教师模型的输出为结构单元,取代传统蒸馏学习中以单个教师模型输出学习的方式,利用多输出组合成结构单元,更能体现出教师模型的结构化特征,使得学生模型得到更好的指导。CC多输入联系型学习法在上述RKD算法的基础上,为了更好的扩大类间差异,更好的缩小类间距离,CC算法提出了两种采样方法:包括均衡类别采样法和均衡超类别采样法。所谓的均衡类别采样法,即假设每个batch大小为48,则这48个样本分布来自于6个类别,每个类别8个样本,使得整体的学习样本不像RKD算法那样是随机的。

(1)Fitnet:阶段性知识蒸馏

https://arxiv.org/pdf/1412.6550.pdf

FItnet这篇文章首次提出了从教室模型的中间层去进行蒸馏学习,而不仅仅关注教室模型的输出。因为通过中间层的引导,使得学生模型进行了提前学习,使得最终的蒸馏学习效果变得更好。

(2)VID:互信息学习法

https://arxiv.org/pdf/1904.05835.pdf

(3)SP:相似性矩阵学习

https://arxiv.org/pdf/1907.09682.pdf

(4)AT:注意力学习法

https://arxiv.org/pdf/1612.03928.pdf

为了更好的表征神经网络中间层的特征,如上图所示,本文列举了三种不同形式的用于更好抽象的表征中间层特征的新的蒸馏形式。其中VID互信息学习法,将中间层知识蒸馏的最优性能定义为最大化教师和学生网络之间的互信息。那么为什么通过最大化互信息可以使得蒸馏学习变得有效呢?首先作者对互信息做了定义:互信息为[教师模型的熵值] - [已知学生模型的条件下的教师模型熵值]。而我们又有如下常识:当学生模型已知,能够使得教师模型的熵很小,这说明学生模型以及获得了能够恢复教师模型所需要的“压缩”知识,间接说明了此时学生模型已经学习的很好了。而这种情况下也就是说明上述公式中的熵很小,从而使得互信息会很大。作者从这个角度解释了为什么可以通过最大化互信息的方式来进行蒸馏学习。而在SP相似性矩阵学习法中,作者提出了一种新的知识蒸馏形式,该方法是作者观察到相似语义的输入往往会使得神经网络输出相似的激活模式这一现象启发得到的。该知识蒸馏方法被称为保持相似性知识蒸馏(SPKD),该方法使得教师网络中相似(不同)激活的输入样本对,能够在学生网络中产生相同(不同)的激活,从而指导学生网络的学习。而在AT注意力学习法中,作者认为注意力在人类视觉体验中起着至关重要的作用。以图像分类为例,注意力地图展示了学习完成后的网络模型更关注于图像的哪个区域,是网络模型学习成果的体现。本文通过迫使学生模型模仿强大的教师模型的注意力特征图,来显著提高学生模型的性能。为此,本文提出了基于激活注意力地图的蒸馏法。

(5)NST:基于滤波器的知识蒸馏

https://arxiv.org/pdf/1707.01219.pdf

(6)FSP:授之以鱼不如授之以渔

http://openaccess.thecvf.com/content_cvpr_2017/papers/Yim_A_Gift_From_CVPR_2017_paper.pdf

和之前对中间层特征进行直接学习的方式不同,本文提出了一个偏哲学的论点:授之以鱼不如授之以渔。具体来说就是,如上图所示,本文将教师模型网络层与层之间的映射关系作为学生网络学习的目标,而不是像之前提到的直接对教师模型的中间结果进行学习。通过让学生学习这种获得特征的方法,而不是直接学习特征本身,文章的结果显示,这种方式确实有助于提高学生模型的鲁棒性。

(1)AB:激活边界学习

https://arxiv.org/abs/1811.03233.pdf

(2)利用对抗样本进行激活边界学习

https://arxiv.org/abs/1805.05532.pdf

在分类任务中,小模型真正的缺陷更多的在于对边界样本(难例样本)的分类困难。而这真是我们在分类任务中最关心的问题。而教师模型处理边界的能力一定是要优于学生模型的。因此尝试用学生模型学习教师模型的边界分布,这将是蒸馏学习的新思路。本部分列举了两种不同的边界学习方法。AB激活边界学习法,通过最大化边界误差的方式,来引导学生模型学习更强的边界约束能力。利用对抗样本进行边界激活学习的方法,首先定义一个基类并通过基类找到各个类别中的对抗边界,最终通过对抗边界样本进行蒸馏学习。

(1)be your own teacherr

https://arxiv.org/pdf/1905.08094.pdf

(2)强制拉近类内距离:regularzing class-wise

https://arxiv.org/pdf/2003.13964.pdf

(3)类内的鲁棒性学习:Data-Distortion Guided

https://www.researchgate.net/publication/335476911_Data-Distortion_Guided_Self-Distillation_for_Deep_Neural_Networks

由于我们不一定可以在所有任务中都顺利的获取教师模型,有的大模型由于数据的缺失,很难被正常的训练出来。基于这种情况,很多研究者提出了自我学习的策略。简单来说该策略就是自己作为自己的老师,进行自我优化。本部分列举了三种自学习的方式。be your own teacher这篇文章将网络较深部分的知识压缩到较浅部分,也就是说该蒸馏策略的教师模型和学生模型来自与同一个模型,这大大降低了蒸馏学习的复杂度,并且通过增加额外的训练检测模型,在不增加前向推理时间的前提下提升了自我学习的能力。举例来说,如上图中以resnet50为例,在每个block之后都接出一个bottleneck作为隐藏层的监督输出模块,并接出一个全连接层作为每个子模块的子分类器。每个子分类器都作为一个小的学生模型,其对应的教师模型为主分类层的输出,最终实现自我蒸馏的学习。另外两篇文章的思路主要从同一个类内的样本出发进行自我学习。其中强制拉近类内距离这篇文章,在训练的过程中,首先我们会从数据迭代器中提取batch1大小的数据,同时选出和当前batch1中类别相同的样本形成batch2,并将两个batch的样本组合起来进行联合训练。具体来说就是每一此计算loss时,从batch1和batch2中各挑选出一个同类样本,在loss计算中尽可能的是的这两个同类样本的输出分布一致,这种方式是一种广义上的自我学习的策略,且这种训练方式能够强制减小类内的差异,且可以利用这种方式减小某些过度自信的异常值的预测。其中增强类内鲁棒性这篇文章,也是从对同一个类别的样本进行联合学习,具体操作如下:对输入batch中每个图片利用不同的数据增强方式增强层两份输入,这两份输入的标签为同一个类别,将两份输入特征concat之后通过卷积层提取全局特征,并将得到的特征向量进行对应的切分,在训练过程中通过最小化切分后特征向量间的差异,从而增强同一个类内的多样性提升鲁棒性,该过程也可以被认为是自己和自己学习。

(1)DML:互相学习

https://arxiv.org/pdf/1706.00384.pdf

(2)知识嫁接

https://arxiv.org/pdf/2001.05868.pdf

与自学习类似的是互相学习策略中不存在教师模型,与自学习不同的是互相学习的方式通用是多个模型之间的学习,而自学习仅仅只有一个模型。其中DML:互相学习这篇文章就是一种典型的互相学习的方式,DML是在训练过程中,几个需要反向传播的待训学生网络协同学习,互相传递知识。每个互相学习的网络都有一个标准的分类Loss和互学习Loss,其中互学习Loss是一个KL散度。 具体而言,两个网络的softmax输出为p1,p2.则互学习的意义在于,对于Net1(Net2亦然),对了提高其泛化能力,使用Net2的p2作为一种后验概率,然后最小化p1,p2的KL散度。而知识蒸馏这篇文章更像是一种广义上的互相学习方法,该文章的主要是思想是并行地训练多个网络,对所有网络的参数进行重要性排序,并另一个并行网络中的更有效的权重替换到当前网络的不重要权重的位置,在训练过程中通过这种重要性权重的互相替换实现互相学习。

(1)GAN对抗学习

https://arxiv.org/pdf/1709.00513.pdf

(2)无监督对抗学习

https://arxiv.org/pdf/1904.01186.pdf

本部分主要列举了两种利用GAN网络进行蒸馏的文章。GAN对抗学习这篇文章就是典型的利用生成对抗网络的例子,具体来说,学生网络作为生成器,生成对应的输出结果,而教师网络用来表征GT信息,而鉴别器主要被用来鉴别学生网络的输出和教师网络的输出,最终学习的目的就是是的学生网络能够欺骗鉴别器,是的鉴别起无法区分出学生网络和教师网络的输出。最终实现学生网络学习到了教师网络的输出特征和分布。而无监督对抗学习这篇文章的出发点有点不一样,由于一些实践问题(如隐私、法律等问题),给定深度网络的训练数据往往不可用,除了一些接口之外,给定网络的架构也是未知的。基于此,本文提出了一种利用生成对抗网络训练高效深度神经网络的新框架。讲预先训练好的教师网络看作一个固定的鉴别器,利用该鉴别器产生的训练样本可以得到最大的鉴别结果。然后,利用生成的数据和教师网络,同时训练出模型尺寸较小、计算复杂度较低的高效网络。