论文《Distilling the Knowledge in a Neural Network》
* 源码以Github为准
1. 数据集
本文使用 fashion_mnist数据集，输入图像大小为28*28，共分为10类。
通过tensoflow加载数据，并对label进行 one hot编码。
import tensorflow as tf
from tensorflow import keras
import numpy as np
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
train_images = train_images/255
test_images = test_images/255
train_labels = tf.one_hot(train_labels, depth=10)
test_labels = tf.one_hot(test_labels, depth=10)
2. 教师模型
本文中使用一个 4层MLP来作为教师模型。
训练过程中，模型最后使用softmax层来计算损失值。
训练结束后，更改最后的softmax层，以便生成软标签，其中T=2。同时，为了防止误操作，将教师模型冻结。
需要注意的是，虽然更改后教师模型不再进行训练，但仍需要使用compile函数进行配置，否则无法调用predict函数。
inputs = keras.layers.Input(shape=(28,28))
x = keras.layers.Flatten()(inputs)
x = keras.layers.Dense(128, activation='relu')(x)
x = keras.layers.Dense(128, activation='relu')(x)
x = keras.layers.Dense(128, activation='relu')(x)
x = keras.layers.Dense(10)(x)
outputs = keras.layers.Softmax()(x)
t_model = keras.Model(inputs, outputs)
t_model.summary()
callback = [keras.callbacks.EarlyStopping(patience=10 ,restore_best_weights=True)]
t_model.compile(optimizer='adam',
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
metrics=['accuracy'])
t_model.fit(train_images, train_labels, epochs=500, validation_data=(test_images, test_labels),callbacks=callback)
x = t_model.get_layer(index=-2).output
outputs = keras.layers.Softmax()(x/3)
Teacher_model = keras.Model(t_model.input, outputs)
Teacher_model.summary()
Teacher_model.trainable = False
Teacher_model.compile(optimizer='adam',
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
metrics=['accuracy'])
3. 学生模型
本文使用一个 2层MLP作为学生模型。
学生模型构建完成后不进行训练，在后续的蒸馏过程中进行训练。
需要注意的是，学生模型最后一层 不加Softmax层。
inputs = keras.layers.Input(shape=(28,28))
x = keras.layers.Flatten()(inputs)
x = keras.layers.Dense(128, activation='relu')(x)
outputs = keras.layers.Dense(10)(x)
Student_model = keras.Model(inputs, outputs)
Student_model.summary()
4. 知识蒸馏过程
学生模型进行蒸馏时，损失函数包括两部分：
- Loss1：学生模型softmax输出值与真实标签的之间的损失（交叉熵）；
- *Loss2：学生模型软化后的softmax输出值（T=2）与教师模型生成的软标签之间的损失（KL散度）。
则， Loss = 0.1Loss1 + 0.9Loss2。
本文通过重写Model类来实现。
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)
蒸馏过程加入早停止机制，监视val_loss。
distill = Distilling(Student_model, Teacher_model, 2, 0.9)
distill.compile(optimizer='adam',
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False))
callback = [keras.callbacks.EarlyStopping(patience=20, restore_best_weights=True)]
distill.fit(train_images, train_labels, epochs=500, validation_data=(test_images, test_labels), callbacks=callback)
5. 实验结果
为了验证结果，本文独立训练学生模型（加入Softmax层），与使用知识蒸馏训练的学生模型进行对比。
实验结果如下：
- 教师模型准确度 0.8682
- 学生模型准确度 0.8365 （知识蒸馏）
- *学生模型准确度 0.8302 （独立训练）
这表明，知识蒸馏方法确实有效。
欢迎评论留言讨论交流！
Original: https://blog.csdn.net/For_learning/article/details/117304450
Author: 今生有幸.~
Title: 知识蒸馏 示例代码实现及下载
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/690764/
转载文章受原作者版权保护。转载请注明原作者出处！