def train_step(self, data):
x, y = data
with tf.GradientTape() as tape:
logits_pre = self.student_model(x)
t_logits_pre = self.teacher_model(x, training=False)
clf_loss_value = self.clf_loss(y, tf.math.softmax(logits_pre))
kd_loss_value = self.kd_loss(tf.math.softmax(t_logits_pre/self.T), tf.math.softmax(logits_pre/self.T))
sum_loss_value = self.alpha * clf_loss_value + (1-self.alpha) * kd_loss_value
self.optimizer.minimize(sum_loss_value, self.student_model.trainable_variables, tape=tape)
self.compiled_metrics.update_state(y, tf.math.softmax(logits_pre))
self.sum_loss_tracker.update_state(sum_loss_value)
self.clf_loss_tracker.update_state(clf_loss_value)
self.kd_loss_tracker.update_state(kd_loss_value)
return {m.name: m.result() for m in self.metrics}
上述训练过程可以看成模型蒸馏其实就是计算学生和老师的loss,让学生向着老师loss的大方向和自己的小方向做更新。所以第一步当然是得到两个不一样的loss,只不过老师模型就不需要更新了。
loss来自hard target的loss和soft target的loss,权值往往需要softtarget的权值更大些
clf_loss:对应的是students模型计算的损失
kd_loss:对应的是teachers模型输出的损失,T为温度,作用是使模型的训练困难化
teachers模型的参数不发生变化
值得注意的点是:logits_pre 和 t_logits_pre 都是特征,没有经过softmax,所以并不是先softmax再/T,而是先/T在进行softmax,这对最终形成的loss的方差可以大大减小,因为softmax是x的指数函数与sum的指数函数的比值。