CNN_模型压缩加速&移动端部署_蒸馏算法
总结:离线蒸馏重点在使用小网络获得大网络的能力,在线蒸馏在于提升模型的精度和训练时间。
现有的加速库有:Intel的Neural Network Distiller和腾讯的PocketFlow。
一 、离线蒸馏
2014 年Hinton 提出了「dark knowledge」的概念。Dark knowledge,顾名思义,就是隐藏在深度学习表面上所展现出来的网络结构、节点之间的连接权重、网络的输出这些看得到的数据之下的知识。在他看来,这种被称为「暗知识」的东西才是深度学习本质上学到的知识(或许这也是寻求深度学习可解释性的一种途径)。如果能够找到一种途径,使得我们能够获取这种知识,并且将其包装成一种先验概率,迁移到更小的模型中去,是否能够提升小模型的性能呢?有人认为,这也可以看作是在迁移学习的框架下,将大规模网络视作信息充分的 source domain,将小规模网络视作需要大规模网络补充信息的 target domain,而我们提取到的 dark knowledge 则是两个任务之间的 common knowledge。
我们在绝大多数的预测任务的深度学习网络中,都会使用 softmax layer 为大量的标签分配概率分布。然而这种处理方式存在一个负作用:与正确标签相比,模型为所有的误标签都分配了很小的概率;然而实际上对于不同的错误标签,其被分配的概率仍然可能存在数个量级的悬殊差距。例如:在图片分类问题中,我们要将图片分成猫、狗、老虎三类。在一次训练中,我们给三类分配的概率分别为 [0.0010, 0.0001, 0.9989],从而最终得到 [0,0,1] 的 one-hot 编码作为分类结果(即 hard-target),我们认为图片所代表的是一只老虎。然而,softmax 函数输出的概率往往包含着类别之间潜在的相关性。在这个例子中,我们可以看到,图片可能是猫的概率比图片是狗的概率更接近图片是老虎的概率,这说明猫和老虎之间存在的内在联系更加强大。由于在宏观上由于这些概率都很小,这一部分的知识很容易在训练过程中淹没,这无疑是浪费了重要的可以用于将大规模网络的知识迁移到小规模网络中去的宝贵先验概率。
为了充分利用这种类类别之间的相关性,我们需要通过某种方式去改变概率分布,使其更加平缓。Hinton 向 softmax 函数添加了一点「佐料」——参数「T, 温度」
其中,z 为每一个类别输入的 logit。式中 T=1 时,退化成传统的 softmax;T无穷大时,结果趋近于 1/C,即所有类别上的概率趋近于相等。T>1 时,我们就能获得 soft target label。通过提高 T,softmax层的映射曲线更加平缓,因而实例的概率映射将更为集中,便使得目标更加地「soft」。
有了这个 distillation 的内核,Hinton 按照以下的步骤对大规模网络进行「蒸馏」:
1. 训练大模型:先用 hard target(类似于 [0,0,1] 的 one-hot 编码)的样本训练。
2. 计算 soft target:利用训练好的大模型来计算 soft target 。也就是大模型「软化后」再经过 softmax 的输出。
3. 重新创建一个小的网络,该网络最后有两个 loss,一个是 hard loss,即传统的 softmax loss,使用 one-hot label;另外一个是 soft loss,即 T>1 的 softmax loss,使用我们第二步保存下来的 soft target label。
整体的 loss 如下式:
其中
为第二步输出的 soft label。
用「软化」训练集训练小模型。训练小模型时 T 不变仍然较大,训练完之后 T 改为1。
4. 预测时,将训练好的小模型按常规方式使用。
二、在线蒸馏
来自谷歌、谷歌大脑、DeepMind的研究人员,包括Geoffrey Hinton等人,ICLR 2018的论文“Large scale distributed neural network training through online distillation”。
在实践中,codistillation比精心调参的offline distillation得到的质量提升是次要的,更有趣的研究方向是将codistillation作为一种分布式训练算法来研究。目前,分布式随机梯度下降(SGD),包括其同步和异步形式(Chen et al.,2016)是在多个互联机器上进行大规模神经网络训练的主要算法。但是,随着机器数量的增加,训练一个高质量模型所需时间的改善程度会降低,直到继续增加机器却无法进一步缩短训练时间。基础架构的限制,以及优化上的障碍,一起限制了分布式 minibatch SGD的可扩展性。
一旦我们达到了向分布式SGD添加worker的限制,我们就可以使用额外的机器来训练模型的另一个副本,并创建一个集成(ensemble)以提高准确性(或通过以更少的step训练ensemble中的成员来提高训练时的精度)。ensemble能够做出更稳定和可重复的预测,这在实际应用中很有用。但是,ensemble增加了测试时的成本,可能会影响延迟或其他成本限制。
从时间和pipeline复杂性两方面来看,额外的训练成本阻碍了从业者使用ensemble distillation,尽管这种方法基本上总是能够改善结果。在这项新的工作中,我们描述了一个简单的distillation的在线变体,我们称之为codistillation。Codistillation通过向第i个模型的损失函数添加一个项来匹配其他模型的平均预测值,可以并行训练n个模型的副本。
通过大规模实验,我们发现,与分布式SGD相比,通过允许有效利用更多计算资源, codistillation提高了准确性并加快了训练速度,甚至加速效果超过了给SGD方法添加更多worker。
流程:
1、所有模型使用相同的架构;
-
2、使用相同的数据集来训练所有模型;
-
3、在任何模型完全收敛之前使用训练期间的distillation loss。
相关的另一个尝试《Born Again Neural Networks》(论文地址:https://arxiv.org/abs/1805.04770)
我们观察到,由知识蒸馏引起的梯度可以分解为两项:含有错误输出信息的暗知识(DK)项和标注真值项,后者对应使用真实标签获得原始梯度的简单尺度缩放。我们将第二个术语解释为基于教师模型对重要样本的最大置信度,使用每个样本的重要性权重和对应的真实标签进行训练。这说明了 KD 如何在没有暗知识的情况下改进学生模型。
此外,我们还探讨了 Densenet 教师模型提出的目标函数能否用于改进 ResNet 这种更简单的架构,使其更接近最优准确度。我们构建了复杂性与教师模型相当的 Wide-ResNet(Zagoruyko & Komodakis,2016b)和 Bottleneck-ResNet(He 等,2016 b)两个学生模型,并证明了这些 BAN-ResNet 性能超过了其 DenseNet 教师模型。类似地,我们从 Wide-ResNet 教师模型中训练 DenseNet 学生模型,前者大大优于标准的 ResNet。因此,我们证明了较弱的教师模型仍然可以提升学生模型的性能,KD 无需与强大的教师模型一起使用。
图 1:BAN 训练过程的图形表示:第一步,从标签 Y 训练教师模型 T。然后,在每个连续的步骤中,从不同的随机种子初始化有相同架构的新模型,并且在前一学生模型的监督下训练这些模型。在该过程结束时,通过多代学生模型的集成可获得额外的性能提升。
参考链接:https://zhuanlan.zhihu.com/p/35698635