pytorch中kl散度蒸馏必须用f.log_softmax(student_logits/t)和f.softmax(teacher_logits/t)配对,顺序不可反,且需乘t²;教师前向须用torch.no_grad()包裹并调用eval(),学生保持train()模式;总损失为alpha×kl_loss+(1−alpha)×ce_loss,alpha通常取0.1–0.4。

PyTorch里KL散度计算软标签必须用 F.log_softmax 和 F.softmax 配对
直接拿教师模型输出的 logits 做 kl_div 会报错或结果发散——因为 PyTorch 的 F.kl_div 要求第一个输入是 log-probabilities(即 log-softmax),第二个是 probabilities(即 softmax),且默认 reduction 是 'batchmean'。
常见错误写法:F.kl_div(teacher_logits, student_logits, ...),这既没归一化也没对齐分布类型。
- 正确做法:教师侧用
F.softmax(teacher_logits / T, dim=1)得到软概率,学生侧用F.log_softmax(student_logits / T, dim=1) - 温度系数
T通常取 3–7;T=1时退化为硬标签交叉熵,梯度变弱 - 务必保持两者的
dim一致,且kl_div输入顺序不能反(log-prob 在前,prob 在后) - 示例片段:
import torch.nn.functional as F<br>kl_loss = F.kl_div(<br> F.log_softmax(student_logits / T, dim=1),<br> F.softmax(teacher_logits / T, dim=1),<br> reduction='batchmean'<br>) * (T ** 2)
注意最后乘T ** 2才能保证梯度尺度合理
教师-学生网络联合训练时,torch.no_grad() 只加在教师前向,不冻结其参数
教师模型在蒸馏中只提供监督信号,本身不参与反向传播,但它的权重仍需保持固定——不是靠 requires_grad=False,而是靠上下文管理器屏蔽梯度。
- 错误操作:给教师模型调用
model.eval()后就以为万事大吉;其实eval()只影响 Dropout/BatchNorm,不影响梯度流 - 必须显式包裹教师前向:
with torch.no_grad():<br> teacher_logits = teacher_model(x)
- 学生模型必须保持
train()模式,否则 BatchNorm 统计量不更新、Dropout 不生效 - 如果教师模型含 BatchNorm 层(如 ResNet),即使加了
no_grad,也建议提前调用teacher_model.eval()避免统计量被意外更新
蒸馏损失和原始任务损失要按比例混合,alpha 通常设在 0.1–0.4 区间
纯 KL 散度收敛快但泛化差,纯交叉熵收敛慢但最终精度高。二者不是简单相加,比例选错会导致学生模型完全学偏或蒸馏失效。
- 典型组合:
total_loss = alpha * kl_loss + (1 - alpha) * ce_loss,其中ce_loss是学生对真实标签的交叉熵 -
alpha过大(>0.5):学生过度拟合教师输出,尤其在教师有误判时会继承错误 -
alpha过小( - 更鲁棒的做法是 warmup:前几个 epoch 用
alpha=0先训稳学生,再线性升到目标值
BatchNorm 层在学生模型中容易因小 batch size 导致蒸馏不稳定
当学生用轻量网络(如 MobileNetV2)且 batch size 较小(≤32)时,BN 统计不准,导致软标签匹配失真——KL 损失震荡,准确率卡在 70% 上不去很可能是这个原因。
- 验证方法:把学生模型中所有
BatchNorm2d替换为nn.Identity(),若 KL 损失变平滑、精度上升,基本可定位 - 临时解法:改用
SyncBatchNorm(多卡)或增大 effective batch size(梯度累积) - 长期建议:学生模型优先选 GroupNorm 或 LayerNorm 结构(如 Deformable DETR 中的蒸馏实践)
- 别忽略数据增强一致性:教师和学生必须用完全相同的 transform,否则软标签与输入不对齐
实际部署时最容易被跳过的点:教师模型输出的 logits 必须和学生在同一设备(GPU/CPU),且 dtype 一致(通常是 float32)。混用 float16 教师 + float32 学生会导致 KL 计算溢出或 nan,而 PyTorch 默认不报错,只让 loss 突然飙高。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











