好饭不怕晚,这又是一篇很晚很晚的饭,10年前的论文咱们今天才来看。不过,我们始终相信经过时间洗礼和沉淀下来的东西,才是真正有价值的。
从今天开始,将会以 DeepSeek 为主线陆续来写作一系列与大模型相关的技术文章,依旧会不急不躁尽最大努力把一些细节给介绍清楚。同时,为了让大家体会和感触到原作者想法,后续我们也会想之前解析 Transformer 这篇论文一样,会把其中关键的原文也给摘录出来。
因为 DeepSeek 后续很多模型都是蒸馏(Distillation)而来,所以有必要把先把这项技术的细节给大家交代清楚。
最后,经过本文多轮实验对比,在 20 个 epoch 后在教师模型在 FashionMNIST 测试集上的准确率是0.887;不使用蒸馏的学生模型直接训练 5 个 epoch 后的准确率为 0.838;保持同样参数下使用蒸馏方式的学生模型 5 个 epoch 后的准确率为0.865。
并且,如果在构建损失函数的时候,把 soft logits 和 logtis 之间的交叉熵换成 KL 散度,那么同样的参数下蒸馏后的准确率还能达到 0.871左右。
关键字:模型迁移、模型蒸馏、DeepSeek、交叉熵损失、信息熵
1. 模型蒸馏与迁移#
一提到模型蒸馏另外一个能与之呼应的词大概就是模型迁移,也就是常说的迁移学习(Transfer Learning)。那模型蒸馏和模型迁移有什么区别呢?简单来讲,模型迁移是将模型的 A 的部分或全部能力迁移到模型 B 中,而迁移的过程本质上就是将用 A 中的权重参数来重新初始化 B 中对应的权重参数,这也就暗含着迁移部分的网络结构是一样;而模型蒸馏则更像是拿复杂模型(Cumbersome /ˈkʌmbərsəm/ Model )A 来作为一个超牛逼的老师,然后来指导学生模型 B 的训练使得它能够模仿输出和复杂模型类似的结果,两者的网络结构可以是截然不同的。
如图1所示,可以表示模型迁移与模型蒸馏的区别。在图1上,黑色加粗线条部分的网络结构便是从模型 A 迁移到模型 B 的部分;而图1下侧,蒸馏过程则是模型 B 借助模型 A 的输出内容(能力)来指导自己学习,和模型结构没有关系。
更多关于模型迁移部分的内容可以参见深度学习栏目「第5.4节 迁移学习教程:预训练模型微调与迁移训练」内容。
A conceptual block that may have prevented more investigation of this very promising approach is that we tend to identify the knowledge in a trained model with the learned parameter values and this makes it hard to see how we can change the form of the model but keep the same knowledge.
上面这段话就很好的诠释了我们上面提到的模型迁移与模型蒸馏的差异点。因为从传统的观点来看,模型的能力就是等价与模型的权重参数,而能力的迁移也就暗含着权重参数的迁移,而在模型蒸馏中却不是这样。因此上面最后一句话说到,我们有点难以想象如何在改变模型结构以后还能保持模型能力的迁移。
当然,对于模型迁移和模型蒸馏来说还有一点区别,通常来说在模型迁移中两个模型的需要完成的任务场景并不相同,例如将使用 ImageNet 训练得到的网络模型的部分权重参数迁移到其它模型中,然后该模型基于此来构建后续的任务场景;而模型蒸馏则是两个模型的所要完成的任务是一样的,只是说学生模型的结构更简单了。
2. 模型蒸馏的动机#
我们一贯强调,任何一个新技术的出现,要么是为了解决新问题,要么就是为了解决已有技术所存在的弊端,而模型蒸馏技术的出现便是属于后者。首先对于模型蒸馏这个词(Distillation)来说,它是由深度学习领域中的超级大佬杰弗里·辛顿(Geoffrey Hinton /ˈdʒɛfri ˈhɪntən/)于2015年所提出 [1],但是对于这项技术的思想在他之前的 Caruana 等人就已经提出过 [2],Hinton 只是在此基础上进行了改进。
Once the cumbersome model has been trained, we can then use a different kind of training, which we call “distillation” to transfer the knowledge from the cumbersome model to a small model that is more suitable for deployment.
通常,在深度学习中我们为了提高模型的效果可以采用多个模型集成(Ensemble)或者以加大模型复杂度的方式进行,例如2015年谷歌所提出的 「4.8 GoogLeNet网络:Inception 结构与建模思路」 就采用了模型集成的方式。但是,这由此带来的一个问题就是模型在推理的时候所耗费的资源太高,且不利于线上部署。
A very simple way to improve the performance of almost any machine learning algorithm is to train many different models on the same data and then to average their predictions. Unfortunately, making predictions using a whole ensemble of models is cumbersome and may be too computationally expensive to allow deployment to a large number of users, especially if the individual models are large neural nets.
因此,Hinton 等人在 Caruana 的基础之上,提出了一种新的模型压缩方法,也就是蒸馏,来将一个复杂模型的能力转移到一个结构相对简单的模型上,但是还能最大程度地保留复杂模型的精度。
Caruana and his collaborators have shown that it is possible to compress the knowledge in an ensemble into a single model which is much easer to deploy and we develop this approach further using a different compression technique.
3. 模型蒸馏#
3.1 模型蒸馏思想#
那说了这么多,模型蒸馏的原理到底是怎么样的呢?在回答这个问题之前我们先来回顾一下,对于任意一个模型来说(以分类模型为例),模型最后输出的原始结果是什么?
对,就是当前样本属于每个类别下的概率!也就是一个概率分布。
假定对于一个三分类模型来说,它输出的概率分布为 $[0.35,0.28,0.37]$,那么它对应的预测标签便为 $[0,0,1]$。 此时我们能从预测标签中得到什么信息?很单一,除了知道这个样本属于最后一个类别以外并不知道任何信息。
但是根据这个概率分布我们能得到什么呢?①样本属于最后一个类别;②模型对于该样本的判别能力并不是特别强;③模型对于这类样本来说似乎容易将其与第1个类别混淆;④ 每个类别间的相关性。……
看到没有,只要我们想要挖掘,我们便可以从这个概率分布得到很多很多的信息,而这比起一个单独 one-hot 形式的硬标签(Hard Targets)所提供的信息要多得多。同时,这些错误答案之间的相对概率还能够告诉我们模型如何对新数据进行泛化的,反映出了模型的泛化方式。也就是说,错误答案的概率差异反映了模型在不同情况下的推理偏向或者不确定性。
例如,论文中举了这样一个例子。假设有一张 BMW 车的图片,模型只有很小的概率会将其误判为垃圾车。但即使是被误判为垃圾车,概率仍然比误判为胡萝卜的概率大得多。这说明,虽然模型会犯错,但它的错误判断有其偏好,比如将 BMW 错误地分类为垃圾车,比分类为胡萝卜更有可能性。
既然我们能从概率分布中得出这么多有用的信息以及它是如何对新样本进行泛化的,那怎么样充分利用这一点来让复杂模型来指导简单模型学习,让简单模型也能模仿复杂模型输出这样的概率分布,即进行模型转移呢?
一个显而易见的方法就是将复杂模型输出的概率分布作为简单模型训练时的软标签(Soft Targets),以此来让简单模型获得复杂模型的泛化能力。当然,如果复杂模型也是由一系列的模型集成而来,那么此时对于这个复杂模型的输出概率来说,可以取这些一些列模型概率输出的算术平均或几何平均来作为软标签。
An obvious way to transfer the generalization ability of the cumbersome model to a small model is to use the class probabilities produced by the cumbersome model as “soft targets” for training the small model.
3.2 模型蒸馏原理#
在此之前,Caruana 等人通过以网络的输出 logits 来作为软标签让简单模型进行学习,并通过均方误差来最小化整个损失函数,但是后续我们会证明这其实就是模型蒸馏中的一种特殊情况;而 Hinton 等人改进的蒸馏方法则是通过如下方式来得到软标签(模型蒸馏时的目标函数有两部分,即分别对软标签和硬标签的拟合)
$$ q_i=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)}\tag{1} $$其中 $z_i$ 表示原始输出的 logits,$T$ 表示温度。可以看出,当 $T=1$ 时对应的便是正常的 Softmax 操作。
Caruana and his collaborators circumvent this problem by using the logits (the inputs to the final softmax) rather than the probabilities produced by the softmax as the targets for learning the small model and they minimize the squared difference between the logits produced by the cumbersome model and the logits produced by the small model.
Our more general solution, called “distillation”, is to raise the temperature of the final softmax until the cumbersome model produces a suitably soft set of targets. We then use the same high temperature when training the small model to match these soft targets. We show later that matching the logits of the cumbersome model is actually a special case of distillation.
那 Hinton 等人为什么要采用式(1)中的过程来进行处理呢?
3.3 蒸馏中的损失函数#
在回答上面这个问题以前,我们先来看一下在模型蒸馏的过程中所用到损失函数。在模型蒸馏中损失函数一共包含两个部分:
① 蒸馏模型的 logits 经式(1)处理后($T=1$),再同正确标签(硬标签)进行交叉熵损失计算,也就是我们分类中通常使用的方式;
② 蒸馏模型的 logits 和复杂模型的 logits 分别经式(1)处理后(使用同一个$T$),再进行交叉熵损失计算。
最后,再将两部分损失加权平均作为蒸馏模型训练时的整体损失函数。
Typically, the small model cannot exactly match the soft targets and erring in the direction of the correct answer turns out to be helpful.
具体地,设复杂模型第 $i$ 个样本输出的 logits 为 $c_i$ ,蒸馏模型输出的 logits 为 $d_i$ ,样本的正确标签为 $y_i$ ,则整个损失函数的计算过程为
根据式(1)计算蒸馏模型与正确标签之间的交叉熵损失
$$ \begin{aligned} \hat{y}_i=\frac{\exp(d_i/1)}{\sum_j\exp(d_j/1)},\quad L_{\text{hard}}=-\sum_j y_j\log(\hat{y}_j) \end{aligned}\tag{2} $$根据式(1)计算蒸馏模型与软标签之间的交叉熵损失(提醒:这里 $p_i$ 是标签)
$$ \begin{aligned} q_i=\frac{\exp(d_i/T)}{\sum_j\exp(d_j/T)},\quad p_i=\frac{\exp(c_i/T)}{\sum_j\exp(c_j/T)},\quad L_{\text{soft}}=-\sum_j p_j\log(q_j) \end{aligned}\tag{3} $$整体损失
$$ L_{\text{total}}=\alpha L_{\text{soft}}+(1-\alpha)L_{\text{hard}}\tag{4} $$其中,$\alpha$ 是一个超参数,控制软标签和硬标签的权重。通常,软标签的损失函数会占据较大权重(例如,$\alpha = 0.7$),而硬标签的损失则占据较小的权重(例如,$1 - \alpha = 0.3$)。
we found that a better way is to simply use a weighted average of two different objective functions. The first objective function is the cross entropy with the soft targets and this cross entropy is computed using the same high temperature in the softmax of the distilled model as was used for generating the soft targets from the cumbersome model. The second objective function is the cross entropy with the correct labels. This is computed using exactly the same logits in softmax of the distilled model but at a temperature of 1. We found that the best results were generally obtained by using a condiderably lower weight on the second objective function.
但是这里还有一个问题,由于通过软标签生成的梯度的幅度会随着温度 $T$ 的平方倒数缩小(即与$\frac{1}{T^2}$ 成正比),因此当同时使用硬标签和软标签时,需要将软标签的梯度乘以 $T^2$。
具体地,根据式(3)可知,$L_{\text{soft}}$ 关于 $d_i$ 的梯度为
$$ \frac{\partial L_{\text{soft}}}{\partial d_i}= \frac{\partial L_{\text{soft}}}{\partial d_i/T}\frac{\partial d_i/T}{\partial d_i}=\left(q_i-p_i\right)\frac{1}{T}\tag{5} $$注:式(5)中 $\left(q_i-p_i\right)$ 的推导过程可以参见深度学习栏目「第3.7节 从零实现分类模型:Softmax 回归代码解析」内容。
将式(3)代入式(5)有
$$ \frac{\partial L_{\text{soft}}}{\partial d_i}=\frac{1}{T}\left(\frac{\exp(d_i/T)}{\sum_j\exp(d_j/T)}-\frac{\exp(c_i/T)}{\sum_j\exp(c_j/T)}\right)\tag{6} $$如果 $T$ 的取值相对于 logits 来说较大,那么式(6)将会近似等于
$$ \begin{aligned} \frac{\partial L_{\text{soft}}}{\partial d_i}&\approx\frac{1}{T}\left(\frac{1+d_i/T}{N+\sum_j d_j/T}-\frac{1+c_i/T}{N+\sum_jc_j/T}\right)\\[3ex] &=\frac{1}{T}\left(\frac{T+d_i}{T(N+\sum_jd_j/T)}-\frac{T+c_i}{T(N+\sum_jc_j/T)}\right)\\[3ex] &=\frac{1}{T^2}\left(\frac{T+d_i}{N+\sum_jd_j/T}-\frac{T+c_i}{N+\sum_jc_j/T}\right) \end{aligned} \tag{7} $$其中 $N$ 表示分类数量,且第1行中分子分母中的化简原理是 $e^x$ 的一阶泰勒展开,即 $e^x\approx1+x$ 。
可以看出,随着温度 $T$ 的增大,$L_{\text{soft}}$ 关于 $d_i$ 的梯度会越来越小,以 $T^2$ 的规模进行缩小,因此可以通过乘以 $T^2$ 来确保训练过程中的梯度尺度合适。
最终,式(4)可以改写为
$$ L_{\text{total}}=\alpha T^2 L_{\text{soft}}+(1-\alpha)L_{\text{hard}}\tag{8} $$3.4 高熵有利于模型蒸馏#
在介绍完蒸馏中的损失函数以后,我们再回过头去看为什么要采用式(1)这个问题。
行文至此,这里想问大家一个问题,我们对模型蒸馏的目的是什么?不就是想把复杂的教师模型的能力蒸馏到简单的学生模型中吗?这意味着什么?意味着教师模型一定是一个表现优异的模型,不然蒸馏什么,比谁差吗?
既然如此,那这就蕴含着教师模型的输出分布其实是不那么平滑的。例如,对于 MNIST 这个数据集来说,教师模型的对于每个类别来说其输出分布的置信度可能都非常高,类似 $[0.01,0.96,0.02,...]$。那这会带来什么问题呢?
① 越是不平滑的概率分布,首先其对应的信息熵便越低,提供给学生模型的信息就越少
注:关于信息熵可以理解为一个概率分布所携带的信息量,信息量越大也就表示不确定性越高,对应的信息熵也就越大。
例如对于第3.1节中的三分类概率分布来说,它携带的信息熵便为
$$ H=-(0.35\log0.35+0.28\log0.28+0.37\log0.37)\approx1.58\tag{9} $$假定概率分布变为 $[0.05,0.05,0.9],$此时的信息熵是
$$ H=-(0.05\log0.05+0.05\log0.05+0.9\log0.9)\approx0.57\tag{10} $$从不确定性的角度来看,第一种情况的不确定性远远高于第二种,因此前者的信息熵更高。而具有高信息熵的前提就是概率分布尽可能要平滑。
如图2所示,随着温度的增加,整个分布会更加的趋于平滑,并且信息熵也在递增。
② 还会影响交叉熵的计算,造成计算结果不稳定,因为 $\log x$ 在区间 $(0,+\infty)$ 上的值域为 $(-\infty,+\infty)$,反应到训练过程中就是不同样本间梯度方差大。
假设我们有一个训练样本,真实标签为类别 B,即 $y_{\text{hard}}=[0,1,0]$ 表示其属于类别 A、B 和 C 的概率分别是0、1和 0。同时,软标签 $y_{\text{hard}}=[0.2,0.7,0.1]$ ;前向传播计算得到的网络输出(即预测的概率分布)为 $\hat{y}=[0.4,0.5,0.1]$ 。
在使用交叉熵作为损失函数的背景下,损失函数对网络输出 $\hat{y}$ 的梯度为
$$ \frac{\partial L}{\partial \hat{y}_i}=-\frac{y_i}{\hat{y}_i}\tag{11} $$对于硬标签来说,因为硬标签只关注类别 B,所以梯度只与类别 B 相关
$$ \frac{\partial L_{\text{hard}}}{\partial \hat{y}_1}=-\frac{0}{0.4}=0,\quad\frac{\partial L_{\text{hard}}}{\partial \hat{y}_2}=-\frac{1}{0.5}=-2,\quad \frac{\partial L_{\text{hard}}}{\partial \hat{y}_3}=-\frac{0}{0.1}=0\tag{12} $$对于软标签来说,每个类别的梯度都有影响,因此所有类别的梯度都会被计算
$$ \frac{\partial L_{\text{soft}}}{\partial \hat{y}_1}=-\frac{0.2}{0.4}=-0.5,\quad\frac{\partial L_{\text{soft}}}{\partial \hat{y}_2}=-\frac{0.7}{0.5}=-1.4,\quad \frac{\partial L_{\text{soft}}}{\partial \hat{y}_3}=-\frac{0.1}{0.1}=-1\tag{13} $$通过观察我们可以得出,硬标签梯度非常集中,只关注类别 B,其他类别的梯度为零,这意味着每个训练样本的梯度更新较为极端。这会导致训练过程中梯度变化较大,尤其是在模型预测错误时,梯度更新可能会非常剧烈。软标签由于每个类别的标签都包含一些信息,梯度更加平滑,所有类别的梯度都有贡献。因此,梯度的变化较为平滑,训练过程中不容易出现剧烈的波动,方差较小。
所以,通过式(1)中的方式来增加 Softmax 的温度,可以让教师的输出分布变得更加平滑。这种平滑的输出概率分布能为学生模型提供更多的信息,尤其是错误类别的相对重要性,而不仅仅是预测正确类别的概率。在高温度下,错误类别的概率不会被压得太小,这样学生模型就能学习到更多类别之间的相似性 和区别。
When the soft targets have high entropy, they provide much more information per training case than hard targets and much less variance in the gradient between training cases.
3.5 模型蒸馏中的特殊情况#
在上面第3.2节中我们说到,Caruana 等人通过均方误差来最小化整个损失函数其实就是 Hinton 提出的蒸馏的一种特殊情况,不过为什么呢?
由式(7)可知,假如对于每个样本来说其 logits 均进行了去均值化处理,那么此时就有 $\sum_jd_j=\sum_jc_j=0$,进一步可得
$$ \begin{aligned} \frac{\partial L_{\text{soft}}}{\partial d_i}&\approx\frac{1}{T^2}\left(\frac{T+d_i}{N+0}-\frac{T+c_i}{N+0}\right)=\frac{1}{NT^2}\left(d_i-c_i\right) \end{aligned} \tag{14} $$因此,根据式(14)可知,其原函数便可以是 $d_i$ 和 $c_i$ 之间的均方误差。
We show later that matching the logits of the cumbersome model is actually a special case of distillation.
4. 模型蒸馏实现#
在介绍完成模型蒸馏的原理以后,下面再来通过一个实际的蒸馏示例给大家演示蒸馏的实战过程。整体来看我们需要实现一个 LeNet5 结构为基础的教师模型,然后再实现一个两层的全连接神经网络作为学生模型,同时为了对比我们会训练两个学生模型,一个使用蒸馏,一个直接训练。
4.1 模型结构实现#
下面,我们开始分别实现一个教师模型和两个学生模型,一个实现三个模型,示例代码如下:
1 class LeNet5(nn.Module):
2 def __init__(self, ):
3 super(LeNet5, self).__init__()
4 self.conv = nn.Sequential(
5 nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2),
6 nn.ReLU(),
7 nn.MaxPool2d(2, 2),
8 nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
9 nn.ReLU(),
10 nn.MaxPool2d(2, 2))
11 self.fc = nn.Sequential(
12 nn.Flatten(),
13 nn.Linear(in_features=16 * 5 * 5, out_features=128),
14 nn.ReLU(),
15 nn.Dropout(0.5),
16 nn.Linear(in_features=128, out_features=256),
17 nn.ReLU(),
18 nn.Dropout(0.5),
19 nn.Linear(in_features=256, out_features=10))
20
21 def forward(self, img, labels=None):
22 output = self.conv(img)
23 logits = self.fc(output)
24 if labels is not None:
25 loss_fct = nn.CrossEntropyLoss(reduction='mean')
26 loss = loss_fct(logits, labels)
27 return loss, logits
28 else:
29 return logits整个上述代码便是教师模型的实现过程,其实就是 LeNet5 这里就不做介绍。
其次,实现后续直接进行训练的学生模型,示例代码如下:
1 class StudentWithoutDistilling(nn.Module):
2 def __init__(self, input_node=28 * 28, output_node=10):
3 super(StudentWithoutDistilling, self).__init__()
4 self.net = nn.Sequential(nn.Flatten(),
5 nn.Linear(input_node, 512),nn.ReLU(),
6 nn.Dropout(0.5),nn.Linear(512, output_node))
7
8 def forward(self, img, labels=None):
9 logits = self.net(img)
10 if labels is not None:
11 loss_fct = nn.CrossEntropyLoss(reduction='mean')
12 loss = loss_fct(logits, labels)
13 return loss, logits
14 else:
15 return logits上述代码就是一个简单的两层神经网络,也不做介绍。
最后,实现使用蒸馏技术的学生模型,示例代码如下:
1 class StudentWithDistilling(nn.Module):
2 def __init__(self, input_node=28 * 28, output_node=10, temperature=2., alpha=0.5):
3 super(StudentWithDistilling, self).__init__()
4 self.net = nn.Sequential(nn.Flatten(), nn.Linear(input_node, 512),
5 nn.ReLU(), nn.Dropout(0.5), nn.Linear(512,output_node))
6 self.temperature = temperature
7 self.alpha = alpha
8
9 def forward(self, img, soft_logits=None, hard_labels=None):
10 logits = self.net(img)
11 if hard_labels is not None and soft_logits is not None:
12 loss_fct = nn.CrossEntropyLoss(reduction='mean')
13 loss_soft = loss_fct(logits / self.temperature,
14 F.softmax(soft_logits / self.temperature,dim=-1))
15 # loss_soft = nn.functional.kl_div(F.softmax(logits / self.temperature),
16 # F.softmax(soft_logits / self.temperature), reduction='batchmean')
17 loss_hard = loss_fct(logits, hard_labels)
18 loss = self.alpha * loss_soft * self.temperature ** 2 + (1 - self.alpha) * loss_hard
19 return loss, logits
20 else:
21 return logits在上述代码中,第2行是实例化时需要传入的参数,其中 temperature 和 alpha 表示式(8)中对应的两个超参数。第9行中 soft_logits 和 hard_labels 分别表示教师模型输出的 logits 和 硬标签,形状分别为 [batch_size, num_class] 和 [batch_size];img 为输入这里用到的是 FashionMNIST 数据集,形状为 [batch_size, 1, 28,28]。第11~18 是根式(2)~式(3)和(8)来计算损失值,这里需要注意的一点是 loss_fct 在计算交叉熵的时候会对第1个参数表示的 input 进行 softmax 计算,所以第13行第1个参数我们不再进行这一步,而第2个参数 target 则需要先进行 softmax 操作;其中,第15~16行是将交叉熵替换成 KL 散度,实验发现替换成 KL 散度的效果会更好,这也是现在很多模型在蒸馏时都会使用 KL 散度来进行代替。
4.2 模型训练实现#
由于教师模型和不使用蒸馏的学生模型训练过程同训练一般模型的过程一致,大家直接阅读源码即可所以这里就不再赘述。唯一需要注意的是教师模型在训练过程中,需要将训练完成后的权重参数持久化保存到本地,以便学生模型蒸馏使用。整个训练过程的代码实现过程如下:
1 from models import StudentWithDistilling
2 from models import LeNet5 as TeacherModel
3
4 def train():
5 epoch, lr, batch_size = 5, 0.001, 64
6 temperature, alpha = 1.5, 0.8
7 train_iter, test_iter = load_dataset(batch_size)
8 student_model = StudentWithDistilling(temperature=temperature, alpha=alpha)
9 teacher_model = TeacherModel()
10 teacher_model_state_dict = torch.load('teacher.pt')
11 teacher_model.load_state_dict(teacher_model_state_dict)
12 teacher_model.eval()
13 optimizer = torch.optim.Adam(student_model.parameters(), lr=lr)
14 for epoch in range(epochs):
15 for i, (x, y) in enumerate(train_iter):
16 soft_logits = teacher_model(x)
17 loss, logits = student_model(x, soft_logits, y)
18 optimizer.zero_grad()
19 loss.backward()
20 optimizer.step()
21 if i % 50 == 0:
22 acc = (logits.argmax(1) == y).float().mean()
23 print(f"Epochs[{epoch + 1}/{epochs}]--batch[{i}/{len(train_iter)}]"
24 f"--Acc: {round(acc.item(), 4)}--loss: {round(loss.item(), 4)}")
25 test_acc = evaluate(test_iter, student_model)在上述代码中,第8~9行分别用于实例化学生模型和教师模型。第10~12行是载入之前持久化保存的教师模型权重参数,并将其赋值给当前的实例化模型 teacher_model 并设置为推理状态。第16~17行分别是根据当前输入计算得到教师模型对应的 soft_logits 和损失及最后的 logits。第25行则是用于测试学生模型在测试集上的准确率。
4.3 实验结果#
经过多轮实验对比,在 20 个 epoch 后在教师模型在 FashionMNIST 测试集上的准确率是0.887;不使用蒸馏的学生模型直接训练 5 个 epoch 后的准确率为 0.838;保持同样参数下使用蒸馏方式的学生模型 5 个 epoch 后的准确率为0.865。
并且,如果在构建损失函数的时候,把 soft logits 和 logtis 之间的交叉熵换成 KL 散度,那么同样的参数下蒸馏后的准确率还能达到 0.871左右。
引用#
[1] Hinton G, Vinyals O, Dean J. Distilling the Knowledge in a Neural Network[J]. stat, 2015, 1050: 9.
[2] C. Buciluaˇ, R. Caruana, and A. Niculescu-Mizil. Model compression. In Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD ’06, pages 535–541, New York, NY, USA, 2006. ACM.