知识蒸馏KL-loss解读

---

前言

本文简单介绍知识蒸馏教师模型与学生模型使用KL loss方法。

一 、KL loss原理

hard label：训练的学生模型结果与真实标签进行交叉熵loss，类似正常网络训练。

soft label：训练的学生网络与已经训练好的教师网络进行KL相对熵求解，可添加系数，如温度，使其更soft。

知乎回答：loss是KL divergence，用来衡量两个分布之间距离。而KL divergence在展开之后，第一项是原始预测分布的熵，由于是已知固定的，可以消去。第二项是 -q log p，叫做cross entropy，就是平时分类训练使用的loss。与标签label不同的是，这里的q是teacher model的预测输出连续概率。而如果进一步假设q p都是基于softmax函数输出的概率的话，求导之后形式就是 q - p。直观理解就是让student model的输出尽量向teacher model的输出概率靠近。

二、logit蒸馏网络图示

三、KL代码

蒸馏模型分类loss代码如下：

import torch
import torch.nn as nn
import numpy as np

loss_f = nn.KLDivLoss()

# 生成网络输出 以及 目标输出
model_student = torch.from_numpy(np.array([[0.1132, 0.5477, 0.3390]])).float() # 假设学生模型输出

model_teacher = torch.from_numpy(np.array([[0.8541, 0.0511, 0.0947]])).float() #假设教师模型输出
label=torch.tensor([0])  # 真实标签
loss_KD = loss_f(model_student, model_teacher)
L=nn.CrossEntropyLoss()
loss_SL=L(model_student,label)
lambda_ ,T=0.6,3  # 分别为设置权重参数，T为温度系数
loss = (1 - lambda_) * loss_SL + lambda_ * T * T * loss_KD  # hint和jeff dean论文

print('\nloss: ', loss)

四、运行结果显示