深度学习 模型压缩之知识蒸馏

知识蒸馏

知识蒸馏是指训练一个参数量较大teacher的模型，用teacher模型去训练参数量较小的student模型，使得student模型具有teacher模型的能力。这里我曾经有个问题，为什么不直接用原有的训练数据集训练student模型，而是用teacher模型的输出作为student的输出的损失？这是因为前者直接训练，标签的信息只包含目标信息（hard targets），如100%是狗。而用后者蒸馏方式标签包含了除了狗之外的信息（soft label）如80%狗+10%猫+10%动物等如下图所示。当看到这，我也有个疑问，这不就是label smooth吗？事实证明，经过这样训练单个模型可以在很大程度上匹配十倍以上集成模型的测试时间性能。在最后softmax层输出，除了正例之外，负标签也带有大量的信息，因此知识蒸馏的方式使得每个样本给student带来的信息量大于传统的训练方式。

蒸馏方式

离线蒸馏

离线蒸馏整个训练过程有两个阶段：

• 大型教师模型是在蒸馏之前首先在一组训练样本上训练
• 教师模型用于提取logit或中间特征形式的知识，然后用于指导蒸馏过程中学生模型的训练

离线蒸馏方法的主要集中于改进知识转移的不同部分，包括知识设计以及用于匹配特征或分布匹配的损失函数。离线的方法主要优点在于它们简单易行。例如，教师模型可以包含使用可能位于不同机器上的不同软件包训练的一组模型。可以提取知识将其存储在缓存中。离线蒸馏方法通常采用单向知识转移和两阶段训练程序。然而，不可避免的是，复杂的高容量教师模型具有很长的训练时间，而离线蒸馏中对学生模型的训练通常在教师模型的指导下是有效的。此外，大型教师和小型学生之间的能力差距始终存在，而学生在很大程度上依赖于教师。

在线蒸馏

尽管离线蒸馏方法简单有效，但离线蒸馏中的一些问题已引起研究界的越来越多的关注。为了克服离线蒸馏的局限性，提出了在线蒸馏以进一步改善学生模型的性能，特别是在没有大容量高性能教师模型的情况下。在在线蒸馏中，教师模型和学生模型同时更新，并且整个知识蒸馏框架是端到端可训练的。在线蒸馏是一种具有高效并行计算功能的单阶段端到端训练方案。然而，现有的在线蒸馏（例如，相互学习）通常不能解决在线蒸馏中的高能力教师，这使得在在线蒸馏中进一步探索教师与学生模型之间的关系成为一个有趣的话题。

自我蒸馏

在自我蒸馏中，教师和学生模型采用相同的网络。这可以视为在线蒸馏的特殊情况。此外，还可以从人类师生学习的角度直观地了解离线，在线和自我蒸馏中。离线蒸馏是指知识渊博的老师向学生传授知识；在线蒸馏是指老师和学生互相学习；自我蒸馏是指学生自己学习知识。而且，就像人类学习一样，这三种蒸馏由于自身的优势可以结合起来互相补充。

蒸馏算法

对抗蒸馏

多教师蒸馏

跨模态蒸馏

图蒸馏

无数据蒸馏

量化蒸馏

深度交互学习（Deep Mutal Learning）

DML也是传统知识蒸馏的扩展，其目标也是将大型模型压缩为小的模型。不同于传统知识蒸馏的单向蒸馏（教师->学生），DML让学生模型相互训练，在整个训练的过程中相互学习，通过这种方式提高模型的性能。DML会让人产生这样的疑问：两个随机初始化的学生网络最初阶段性能都很差的情况，这样相互模型可能会导致性能更差，或者性能停滞不前。针对这样的疑问，会有这样的解释：

• 每个学生主要被传统的有监督学习损失函数影像，这意味着学生网络的性能大体会是增长趋势
• 在监督信号下，所有的网络都会朝着预测正确的label的方向发展，但是不同网络在初始化值不同，他们会学到不同的表征，因此他们对下一类最有可能的概率估计是不同的。
• 在Mutual Learning中，学生群体可以有效汇集下一个最后可能的类别估计，为每个训练实例找到最有可能的类别，同时根据他们互学习对象增加每个学生的后验熵，有助于网络收敛到更平坦的极小值，从而带来更好的泛华能力和鲁棒性。

DML具有的特点是：

• 适合于各种网络架构，由大小网络混合组成的异构网络也可以相互学习
• 效能会随着队列中网络数量的增加而增加，即互学习对象增多的时候，性能会有一定的提升
• 有利于半监督学习，因为其在标记和未标记数据上都激活了模仿损失
• 虽然DML重点是得到某一个有效网络，整个队列中的网络可以作为模型集成的对象进行集成

Demo

环境

keras2.4.3
tensorflow2.4.0

导入相关的包

from keras.datasets import mnist
from keras.layers import *
from keras import Model
from sklearn.metrics import accuracy_score
import numpy as np

加载mnist数据

(data_train,label_train),(data_test,label_test )= mnist.load_data()
data_train = np.expand_dims(data_train,axis=3)
data_test = np.expand_dims(data_test,axis=3)

定义教师模型

def teacher_model():
    input_ = Input(shape=(28,28,1))
    x = Conv2D(32,(3,3),padding = "same")(input_)
    x = Activation("relu")(x)
    print(x)
    x = MaxPool2D((2,2))(x)
    x = Conv2D(64,(3,3),padding= "same")(x)
    x = Activation("relu")(x)
    x = MaxPool2D((2,2))(x)
    x = Conv2D(64,(3,3),padding= "same")(x)
    x = Activation("relu")(x)
    x = MaxPool2D((2,2))(x)
    x = Flatten()(x)
    out = Dense(10,activation = "softmax")(x)
    model = Model(inputs=input_,outputs=out)
    model.compile(loss="sparse_categorical_crossentropy",
                 optimizer="adam",
                 metrics=["accuracy"])
    model.summary()
    return model

训练教师模型

t_model  = teacher_model()
t_model.fit(data_train,label_train,batch_size=64,epochs=2,validation_data=(data_test,label_test))

修改教师模型生成soft label

我觉得这一步是非必要的

T = 3
x = t_model.get_layer(index=-2).output
outputs = Dense(10,activation = "softmax")(x/T)
Teacher_model = Model(t_model.input, outputs)
Teacher_model.summary()
Teacher_model.trainable = False
Teacher_model.compile(loss="sparse_categorical_crossentropy",
                 optimizer="adam",
                 metrics=["accuracy"])

定义学生模型

def student_model():
    input_ = Input(shape=(28,28,1))
    x = Flatten()(input_)
    x = Dense(512,activation="sigmoid")(x)
    out = Dense(10,activation = "softmax")(x)
    model = Model(inputs=input_,outputs=out)
    model.compile(loss="sparse_categorical_crossentropy",
                 optimizer="adam",
                 metrics=["accuracy"])
    model.summary()
    return model
s_model = student_model()

知识蒸馏

import keras
class Distilling(keras.Model):
  def __init__(self, student_model, teacher_model, T, alpha):
    super(Distilling, self).__init__()
    self.student_model = student_model
    self.teacher_model = teacher_model
    self.T = T
    self.alpha = alpha

  def train_step(self, data):
    x, y = data
    softmax = keras.layers.Softmax()
    kld = keras.losses.KLDivergence()
    with tf.GradientTape() as tape:
      logits = self.student_model(x)
      soft_labels = self.teacher_model(x)
      loss_value1 = self.compiled_loss(y, softmax(logits))
      loss_value2 = kld(soft_labels, softmax(logits/self.T))
      loss_value = self.alpha* loss_value2 + (1-self.alpha) * loss_value1
    grads = tape.gradient(loss_value, self.student_model.trainable_weights)
    self.optimizer.apply_gradients(zip(grads, self.student_model.trainable_weights))
    self.compiled_metrics.update_state(y, softmax(logits))
    return {'sum_loss':loss_value, 'loss1': loss_value1, 'loss2':loss_value2, }
  
  def test_step(self, data):
    x, y = data
    softmax = keras.layers.Softmax()
    logits = self.student_model(x)
    loss_value = self.compiled_loss(y, softmax(logits))
    return {'loss':loss_value}

  def call(self, inputs):
    return self.student_model(inputs)

训练

import tensorflow as tf
distill = Distilling(s_model, Teacher_model, 2, 0.9)
distill.compile(loss="sparse_categorical_crossentropy",
                 optimizer="adam",
                 metrics=["accuracy"])
distill.fit(data_train,label_train,batch_size=64,epochs=2,validation_data=(data_test,label_test))

相关论文

可以参考：Awesome-Knowledge-Distillation

总结

知识蒸馏是压缩神经网络的三种主要方法之一。与其他两种强大的模型优化压缩方式剪枝和量化不同，知识蒸馏不直接对网络进行缩减。相反，知识蒸馏训练一个参数量较大的模型“教师模型”来训练一个参数量较小的模型“学生模型”。

对于知识蒸馏，关键是：1）从教师那里提取丰富的知识；2）从教师那里转移知识以指导学生的训练。因此，本文从以下几个方面讨论知识蒸馏的挑战：知识的均等性，蒸馏的类型，师生体系结构的设计以及知识蒸馏的理论基础。

大多数KD方法利用各种知识的组合，包括基于响应的知识，基于特征的知识和基于关系的知识。因此，重要的是要了解每种知识类型的影响，并知道不同种类的知识如何以互补的方式互相帮助。例如，基于响应的知识具有相似的动机来进行标签平滑和模型正则化; 基于特征的知识通常用于模仿教师的中间过程，而基于关系的知识则用于捕获不同样本之间的关系。为此，在统一和互补的框架中对不同类型的知识进行建模仍然是挑战。例如，来自不同提示层的知识可能对学生模型的训练有不同的影响：1）基于响应的知识来自最后一层；2）来自较深的提示/指导层的基于特征的知识可能会遭受过度规范化的困扰。

如何将丰富的知识从老师传授给学生是知识蒸馏的关键一步。通常，现有的蒸馏方法可分为离线蒸馏，在线蒸馏和自蒸馏。离线蒸馏通常用于从复杂的教师模型中转移知识，而教师模型和学生模型在在线蒸馏和自我蒸馏的设置中具有可比性。为了提高知识转移的效率，应进一步研究模型复杂性与现有蒸馏方案或其他新颖蒸馏方案之间的关系。

目前，大多数KD方法都将重点放在新型知识或蒸馏损失函数上，而对师生体系结构的设计研究不足。实际上，除了知识和蒸馏算法之外，教师和学生的结构之间的关系也显着影响知识蒸馏的性能。例如，一方面，最近的一些研究发现，由于教师模型和学生模型之间的模型能力差距，学生模型无法从某些教师模型中学习到很多东西；另一方面，从对神经网络容量的一些早期理论分析来看，浅层网络能够学习与深层神经网络相同的表示。因此，设计有效的学生模型或构建合适的教师模型仍然是知识蒸馏中的难题。

尽管有大量的知识蒸馏方法和应用，但对知识蒸馏的理解（包括理论解释和实证评估）仍然不够。例如，蒸馏可以被视为一种获得特权信息的学习形式。线性教师模型和学生模型的假设使得能够通过蒸馏来研究学生学习特征的理论解释。此外，Cho和Hariharan（2019）对知识蒸馏的功效进行了一些实证评估和分析。但是，仍然很难获得对知识提升的可概括性的深刻理解，尤其是如何衡量知识的质量或师生架构的质量。

参考

1. 知识蒸馏：如何用一个神经网络训练另一个神经网络
2. 深度学习三大谜团：集成、知识蒸馏和自蒸馏
3. 知识蒸馏(Knowledge Distillation) 经典之作
4. 知识蒸馏（Knowledge Distilling），让你的模型轻装上阵
5. Distilling Knowledge in Neural Networks
6. 一文读懂知识蒸馏
7. Knowledge Distillation: A Survey