应根据教师输出选择损失:若为概率分布则用kl散度,需教师softmax后取log、学生log_softmax;若为logits则可用mse但丢失温度平滑效应;kl在小数据或类别不平衡时更稳,但需梯度裁剪。

知识蒸馏中 KL散度损失 和 MSE损失 怎么选?
看教师模型输出的是概率分布还是 logits —— 如果用 torch.nn.KLDivLoss,必须确保学生模型也经过 log_softmax,而教师输出是 softmax;直接拿 raw logits 算 MSE 虽简单,但温度系数 T 没起作用,软标签的平滑效应就丢了。
- 常见错误:把教师的
logits直接喂给KLDivLoss,报错Expected input tensor to have log probabilities - 正确链路:教师
logits / T → softmax → log(),学生logits / T → log_softmax,再进KLDivLoss - 小数据或类别不平衡时,
KLDivLoss通常比MSELoss更稳;但训练初期梯度可能爆炸,建议搭配torch.nn.utils.clip_grad_norm_
PyTorch 里怎么让教师模型不更新梯度?
不是只加 with torch.no_grad(): 就够了 —— 如果教师模型里有 BatchNorm 层,它默认会继续更新 running_mean / running_var,导致蒸馏过程中的统计量漂移,学生学的其实是“动起来的教师”。
- 必须显式调用
teacher.eval(),哪怕你没开 dropout - 如果教师用了
SyncBatchNorm或自定义 BN,还要确认是否在多卡下被torch.nn.parallel.DistributedDataParallel包裹,此时eval()可能不生效 - 更稳妥的做法:在蒸馏循环开头加
teacher.train(False),并手动冻结 BN 参数:for m in teacher.modules(): if isinstance(m, torch.nn.BatchNorm2d): m.track_running_stats = False
temperature 设太高或太低会怎样?
temperature 不是越大越好。设成 20,教师输出几乎变成均匀分布,学生学不到区分性;设成 1,软标签和硬标签没区别,蒸馏失去意义。
- 典型取值在
3~7之间,CV 任务常用4,nlp 任务(如 bert 蒸馏)倾向1~2(因 logits 差异本身已较大) - 如果学生模型比教师小很多(比如 TinyBERT 蒸馏 BERT-base),可以先用
T=5预热 1–2 个 epoch,再降到T=2细调 - 注意:
T同时影响 KL 散度计算和学生输出的梯度尺度,所以学习率得相应调低(经验上除以T²)
蒸馏时要不要保留原始任务 loss?
要,而且不能简单加权平均。学生模型如果只学软标签,会在验证集上出现“高置信、低准确”现象 —— 对教师认为模糊的样本过度自信。
立即学习“Python免费学习笔记(深入)”;
- 标准做法是混合 loss:
total_loss = alpha * ce_loss(student_logits, labels) + (1 - alpha) * kl_loss(student_log_soft, teacher_soft) -
alpha别设成 0.5:初始阶段建议0.2~0.3,等学生 logits 接近教师后再逐步提高到0.5 - 如果下游任务 label 噪声大(比如弱监督数据),可把
ce_loss换成label_smoothing=0.1的版本,避免学生过拟合错误硬标签
事情说清了就结束。真正难的不是调通蒸馏流程,而是发现学生模型在某个子集上始终比教师差 —— 这时候得回过头看教师自己在这个子集上的预测是否稳定,而不是急着改 loss 权重。