知识蒸馏小结

最新推荐文章于 2024-02-20 00:16:56 发布

Dream Algorithm

最新推荐文章于 2024-02-20 00:16:56 发布

阅读量483

点赞数 2

文章标签：深度学习人工智能机器学习神经网络 pytorch Powered by 金山文档

本文链接：https://blog.csdn.net/A2321161581/article/details/129613290

版权

知识蒸馏（knowledge distillation）是模型压缩的一种常用的方法，不同于模型压缩中的剪枝和量化，知识蒸馏是通过构建一个轻量化的小模型，利用性能更好的大模型的监督信息，来训练这个小模型，以期达到更好的性能和精度。

最初的知识蒸馏：

中间层注重纹理细节，深层特征注重抽象语义

如上图所示，教师网络（左侧）的预测输出除以温度参数（Temperature）之后、再做Softmax计算，可以获得软化的概率分布（软目标或软标签），数值介于0~1之间，取值分布较为缓和。Temperature数值越大，分布越缓和；而Temperature数值减小，容易放大错误分类的概率，引入不必要的噪声。针对较困难的分类或检测任务，Temperature通常取1，确保教师网络中正确预测的贡献。硬目标则是样本的真实标注，可以用One-hot矢量表示。Total loss设计为软目标与硬目标所对应的交叉熵的加权平均（表示为KD loss与CE loss），其中软目标交叉熵的加权系数越大，表明迁移诱导越依赖教师网络的贡献，这对训练初期阶段是很有必要的，有助于让学生网络更轻松的鉴别简单样本，但训练后期需要适当减小软目标的比重，让真实标注帮助鉴别困难样本。另外，教师网络的预测精度通常要优于学生网络，而模型容量则无具体限制，且教师网络推理精度越高，越有利于学生网络的学习。

教师网络与学生网络也可以联合训练，此时教师网络的暗知识及学习方式都会影响学生网络的学习，具体如下（式中三项分别为教师网络Softmax输出的交叉熵loss、学生网络Softmax输出的交叉熵loss、以及教师网络数值输出与学生网络Softmax输出的交叉熵loss）：

教师指导学生网络的损失函数代码如下：

def loss_fn_kd(outputs, labels, teacher_outputs, params):
    """
    Compute the knowledge-distillation (KD) loss given outputs, labels.
    "Hyperparameters": temperature and alpha
    NOTE: the KL Divergence for PyTorch comparing the softmaxs of teacher
    and student expects the input tensor to be log probabilities! See Issue #2
    """
    alpha = params.alpha
    T = params.temperature
    KD_loss = nn.KLDivLoss()(F.log_softmax(outputs/T, dim=1),
                             F.softmax(teacher_outputs/T, dim=1)) * (alpha * T * T) + \
                             F.cross_entropy(outputs, labels) * (1. - alpha)
 
    return KD_loss

自蒸馏

自蒸馏（Self Distillation）是三种知识蒸馏学习方法的一种：离线蒸馏（Offline Distillation）、在线蒸馏（Online Distillation）、自蒸馏（Self Distillation）。这三种蒸馏方法区别是：

离线蒸馏（Offline Distillation）: teacher模型是已经预训练好的，已经有很好的效果。再让teacher模型来指导student模型进行学习，完成知识蒸馏。

在线蒸馏（Online Distillation）：teacher模型没有预训练，准备和student模型一同进行训练，在一同训练的过程中来指导student模型进行学习，完成知识蒸馏。

自蒸馏（Self Distillation）：teacher模型和student模型是一个模型，也就是一个模型来指导自己进行学习，完成知识蒸馏。

1、ICCV2019:Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation

知识蒸馏是目前流行的一种压缩方法，其灵感来源于教师对学生的知识转移。其关键策略是将紧凑的学生模型的效果向大型参数化的教师模型近似。然而模型间的知识蒸馏存在以下缺陷：（1）知识迁移的低效性，即学生模型往往并不能充分学习到教师模型的全部知识;(2)如何设计并训练合适的教师模型，现有的蒸馏框架往往需要消耗大量的实验来找到最好的教师模型架构。

如果能将教师网络和学生网络搭建在一个网络中并且能够同时训练，不仅降低了设备要求而且能够提升训练速度，下图是本文提出的自蒸馏策略在训练/推理效率及准确率上的直观对比。