1. 为什么我要在 MNIST 上折腾知识蒸馏MNIST 这个数据集说实话在很多搞深度学习的人眼里就是个“玩具”——28×28 的灰度图10 个类别6 万张训练样本连 LeNet-5 这种上世纪的老古董都能轻松刷到 99% 以上的准确率。我一开始也是这么想的拿它来做知识蒸馏不就是走个过场验证一下代码能不能跑通吗结果我错了而且错得挺离谱。事情是这样的我手头有个实际需求要把一个稍微大一点的模型压缩到一个能在边缘设备上跑的小模型推理延迟要控制在毫秒级参数量最好在 10 万以内。直接拿业务数据做实验成本太高我就想着先在 MNIST 上把知识蒸馏的流程跑通把该踩的坑都踩一遍再迁移到真实场景。这个思路本身没问题问题在于我低估了 MNIST 上做蒸馏的“暗坑密度”。知识蒸馏的核心思想不复杂用一个已经训练好的、性能较强的教师模型去指导一个结构更简单的学生模型训练。教师模型输出的不是硬标签one-hot 的 0 或 1而是软标签——每个类别都有一个概率值比如“这张图有 90% 概率是 37% 概率是 82% 概率是 5”。这些软标签里包含了类别之间的相似性信息也就是所谓的“暗知识”。学生模型学这些软标签理论上能比直接学硬标签获得更多的监督信号。但理论归理论实操归实操。我在 MNIST 上前后跑了十几轮实验踩了四个让我印象深刻的坑每一个都让我重新审视了自己对蒸馏的理解。这篇文章就是把这四个坑掰开揉碎讲清楚包括我是怎么发现的、背后的原因是什么、最后怎么解决的。如果你也打算用 MNIST 或者类似的小数据集做蒸馏实验希望能帮你省下几个小时的调试时间。2. 知识蒸馏在 MNIST 上的整体设计思路2.1 教师-学生框架的基本搭建逻辑我选的是一个经典的教师-学生框架。教师模型用了一个稍微宽一点的 CNN两个卷积层32 和 64 个通道后面接两个全连接层120 和 84 个神经元最后输出 10 类。这个结构参数量大概在 20 万左右在 MNIST 上能轻松跑到 99.3% 以上的测试准确率。学生模型就简单多了一个卷积层16 个通道加一个全连接层32 个神经元参数量压到了 1 万出头。为什么这么选因为我的目标是验证蒸馏在小模型上的效果学生模型必须足够小才能看出蒸馏带来的增益。如果学生模型本身容量就很大那蒸不蒸馏差别不大实验就没有意义了。另外教师模型也不能太复杂否则训练时间太长调试周期会被拉得很长。MNIST 上的 CNN 训练一轮大概几秒钟这个速度很适合快速迭代。损失函数的设计是蒸馏的核心。总损失由两部分组成一部分是学生模型输出和真实硬标签之间的交叉熵损失另一部分是学生模型输出和教师模型软标签之间的蒸馏损失。蒸馏损失用的是 KL 散度但有一个关键细节教师和学生的输出都要先经过温度系数 T的缩放再送进 softmax。公式大概长这样import torch import torch.nn as nn import torch.nn.functional as F class DistillationLoss(nn.Module): def __init__(self, temperature4.0, alpha0.7): super().__init__() self.temperature temperature self.alpha alpha self.ce_loss nn.CrossEntropyLoss() def forward(self, student_logits, teacher_logits, labels): # 硬标签损失 hard_loss self.ce_loss(student_logits, labels) # 软标签损失KL 散度 soft_student F.log_softmax(student_logits / self.temperature, dim1) soft_teacher F.softmax(teacher_logits / self.temperature, dim1) soft_loss F.kl_div(soft_student, soft_teacher, reductionbatchmean) soft_loss soft_loss * (self.temperature ** 2) # 加权求和 total_loss self.alpha * soft_loss (1 - self.alpha) * hard_loss return total_loss这里有两个超参数需要调温度 T 和权重 alpha。T 越大软标签的分布越平滑类别之间的“暗知识”越容易被学生看到但 T 太大也会让分布过于均匀失去区分度。alpha 控制软标签损失和硬标签损失的相对权重。我一开始用的是 T4、alpha0.7这是很多论文里的默认值但在 MNIST 上效果并不理想后面会详细说。2.2 数据加载与预处理中的隐藏陷阱MNIST 的数据加载看起来是最没技术含量的部分但这里恰恰藏着第一个坑。我用的是torchvision.datasets.MNIST标准写法from torchvision import datasets, transforms transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(root./data, trainFalse, downloadTrue, transformtransform)这个写法本身没问题但downloadTrue在某些网络环境下会直接报 404 或者连接超时。我一开始以为是 torchvision 的 bug后来发现是下载源的问题。MNIST 的原始文件托管在几个不同的镜像上torchvision 默认用的那个有时候会抽风。解决办法很简单手动下载四个文件train-images-idx3-ubyte.gz、train-labels-idx1-ubyte.gz、t10k-images-idx3-ubyte.gz、t10k-labels-idx1-ubyte.gz放到./data/MNIST/raw/目录下然后把download改成False。这个坑虽然不涉及蒸馏本身但如果你连数据都加载不了后面的实验根本无从谈起。还有一个容易被忽略的点归一化参数。MNIST 的全局均值是 0.1307标准差是 0.3081这是官方统计出来的。如果你不做归一化或者用错了参数模型收敛会变慢教师模型的最终精度也会受影响。教师模型精度不够高软标签的质量就差蒸馏效果自然打折扣。这个因果关系链很长但每一步都不能马虎。3. 四个坑的详细拆解与解决方案3.1 第一个坑教师模型没训好软标签全是“噪声”我一开始的想法很简单教师模型嘛随便训训就行了反正 MNIST 这么简单。于是我设了 5 个 epoch学习率 0.01也没加学习率衰减训完一看测试准确率 98.7%。当时觉得“还行”就直接拿这个教师模型去生成软标签了。结果学生模型训练的时候loss 震荡得厉害最终准确率只有 97.2%比不用蒸馏直接训还低了 0.5 个百分点。这个结果让我很郁闷。后来我把教师模型的训练日志翻出来仔细看发现 5 个 epoch 的时候教师模型在训练集上的 loss 还在下降测试集准确率也有上升空间。也就是说教师模型根本没收敛。一个没收敛的教师模型输出的软标签质量很差类别之间的概率分布不能反映真实的相似性学生学到的自然是“噪声”。我重新训练了教师模型这次设了 20 个 epoch学习率用余弦退火从 0.01 降到 0.0001并且在第 10 个 epoch 的时候把学习率重启了一次。最终教师模型的测试准确率到了 99.4%训练集上的 loss 也基本平稳了。用这个教师模型重新生成软标签学生模型的准确率直接跳到了 98.6%比不用蒸馏高了 1.4 个百分点。注意教师模型的精度直接决定了蒸馏效果的上限。如果教师模型本身都没训好软标签里的“暗知识”就是错的学生学得越认真错得越离谱。在 MNIST 上教师模型的测试准确率至少要达到 99.2% 以上软标签才有足够的参考价值。这里还有一个细节教师模型在生成软标签的时候一定要用eval 模式并且关闭梯度计算。我见过有人忘了加torch.no_grad()结果显存直接爆了因为教师模型的输出也被挂到了计算图上。正确的写法teacher.eval() with torch.no_grad(): teacher_logits teacher(images)3.2 第二个坑温度系数 T 设得太大软标签“糊”了温度系数 T 是蒸馏里最玄学的超参数。我一开始用的是 T4这是很多教程里的推荐值。但在 MNIST 上T4 的效果并不好。学生模型的收敛速度很慢而且最终准确率比 T2 的时候低了 0.3 个百分点。为了搞清楚原因我把教师模型对某张测试图片的输出打印出来看了看。在 T1 的时候教师模型的输出大概是[0.01, 0.02, 0.91, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01]最大值是 0.91对应数字 2。在 T4 的时候输出变成了[0.08, 0.09, 0.35, 0.08, 0.08, 0.08, 0.08, 0.08, 0.08, 0.08]最大值只有 0.35其他类别的概率都被“抬”起来了。这就是问题所在MNIST 的类别之间区分度很高数字 2 和数字 8 虽然有点像但大部分情况下差异还是很明显的。T4 把分布压得太平滑反而把“2 是正确答案”这个最重要的信息给稀释了。学生模型看到的是一个模棱两可的分布学起来自然费劲。我试了 T1、T2、T3、T4、T6 五组每组跑三次取平均。结果 T2 的时候学生模型表现最好测试准确率 98.6%T1 的时候是 98.3%T4 的时候是 98.3%T6 的时候掉到了 97.9%。所以在这个任务上T2 是一个比较合适的值。温度 T学生测试准确率三次平均收敛 epoch 数198.3%18298.6%15398.4%17498.3%20697.9%25提示温度系数的选择和数据集的类别区分度有关。类别之间越容易混淆比如 CIFAR-10 里的猫和狗T 可以适当大一点类别区分度越高比如 MNISTT 就要小一点。没有一个放之四海而皆准的值必须根据具体任务调。3.3 第三个坑alpha 权重没调好硬标签和软标签“打架”alpha 控制的是软标签损失和硬标签损失的权重。我一开始用 alpha0.7意思是软标签损失占 70%硬标签损失占 30%。这个设置在很多论文里是默认值但在 MNIST 上我发现学生模型的训练 loss 下降得很慢而且最终准确率不如 alpha0.5 的时候。原因其实不难理解。MNIST 的硬标签非常“干净”——每张图片就是一个明确的数字没有标注歧义。硬标签提供的监督信号很强如果硬标签损失的权重太低学生模型就得不到足够的“ ground truth ”信号光靠软标签的“暗知识”是不够的。软标签的价值在于提供类别之间的相似性信息但它不能替代硬标签的准确性。我试了 alpha0.3、0.5、0.7、0.9 四组每组跑三次。结果 alpha0.5 的时候学生模型表现最好测试准确率 98.7%alpha0.3 的时候是 98.5%alpha0.7 的时候是 98.4%alpha0.9 的时候掉到了 98.1%。alpha学生测试准确率三次平均训练 loss 最终值0.398.5%0.120.598.7%0.090.798.4%0.150.998.1%0.21注意alpha 的选择和数据集的质量有关。如果硬标签很干净比如 MNISTalpha 可以设小一点让硬标签损失占更大比重如果硬标签有噪声比如从网上爬下来的数据alpha 就要设大一点更多依赖教师模型的软标签来“纠偏”。3.4 第四个坑学生模型容量太小学不动软标签这个坑是我踩得最深的。我一开始把学生模型压得特别小一个卷积层8 个通道加一个全连接层16 个神经元参数量只有 2000 出头。我想的是“越小越好”这样才能体现蒸馏的价值。结果学生模型的测试准确率死活上不去最高只有 96.8%比不用蒸馏直接训还低了 0.9 个百分点。我一开始以为是蒸馏代码写错了反复检查了损失函数、温度系数、alpha 权重都没问题。后来我把学生模型的容量逐步加大发现准确率跟着往上涨。当学生模型参数量加到 1 万左右的时候蒸馏的效果才真正体现出来测试准确率到了 98.6%。这说明一个问题学生模型必须有一定的容量来“消化”软标签里的信息。如果学生模型太小它的表达能力不足以拟合教师模型输出的复杂分布软标签里的“暗知识”就学不进去。蒸馏不是变魔术它不能把一个 2000 参数的小模型变成 99% 准确率的模型。学生模型和教师模型之间的容量差距不能太大否则蒸馏的增益会被容量瓶颈吃掉。学生模型参数量蒸馏后准确率不用蒸馏准确率蒸馏增益200096.8%97.7%-0.9%500097.9%98.1%-0.2%1000098.6%97.2%1.4%2000098.8%97.5%1.3%提示学生模型的容量选择需要在“压缩率”和“蒸馏增益”之间找平衡。我的经验是学生模型的参数量至少要是教师模型的 1/20蒸馏才能稳定带来正向收益。如果压缩得太狠蒸馏反而会拖后腿。4. 完整实操流程与关键代码实现4.1 教师模型的训练与软标签生成教师模型的训练没什么特别的就是标准的 CNN 训练流程。但有几个细节需要注意学习率用余弦退火训练 epoch 数要足够我用了 20 个并且在训练过程中保存验证集准确率最高的那个 checkpoint。不要用最后一个 epoch 的模型因为最后一个 epoch 可能过拟合了。import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR teacher TeacherNet() optimizer optim.Adam(teacher.parameters(), lr0.01) scheduler CosineAnnealingLR(optimizer, T_max20) best_acc 0.0 for epoch in range(20): teacher.train() for images, labels in train_loader: optimizer.zero_grad() outputs teacher(images) loss F.cross_entropy(outputs, labels) loss.backward() optimizer.step() scheduler.step() # 验证 teacher.eval() correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: outputs teacher(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() acc correct / total if acc best_acc: best_acc acc torch.save(teacher.state_dict(), best_teacher.pth) print(fEpoch {epoch1}, Test Acc: {acc:.4f})软标签的生成很简单就是把训练集重新过一遍教师模型把输出的 logits 保存下来。注意要保存logits不是 softmax 之后的概率值因为蒸馏损失里还要再做一次带温度的 softmax。teacher.load_state_dict(torch.load(best_teacher.pth)) teacher.eval() all_logits [] all_labels [] with torch.no_grad(): for images, labels in train_loader: logits teacher(images) all_logits.append(logits) all_labels.append(labels) all_logits torch.cat(all_logits, dim0) all_labels torch.cat(all_labels, dim0) torch.save({logits: all_logits, labels: all_labels}, teacher_logits.pth)4.2 学生模型的蒸馏训练循环学生模型的训练循环和普通训练差不多区别在于损失函数的计算。我写了一个DistillationLoss类把硬标签损失和软标签损失封装在一起。训练的时候每个 batch 的教师 logits 直接从预先生成的teacher_logits.pth里按索引取这样比每个 batch 都跑一遍教师模型要快得多。student StudentNet() optimizer optim.Adam(student.parameters(), lr0.005) scheduler CosineAnnealingLR(optimizer, T_max30) criterion DistillationLoss(temperature2.0, alpha0.5) teacher_data torch.load(teacher_logits.pth) teacher_logits_all teacher_data[logits] labels_all teacher_data[labels] batch_size 128 num_samples len(labels_all) for epoch in range(30): student.train() indices torch.randperm(num_samples) total_loss 0.0 for i in range(0, num_samples, batch_size): idx indices[i:ibatch_size] images train_dataset.data[idx].float().unsqueeze(1) / 255.0 images (images - 0.1307) / 0.3081 labels labels_all[idx] t_logits teacher_logits_all[idx] optimizer.zero_grad() s_logits student(images) loss criterion(s_logits, t_logits, labels) loss.backward() optimizer.step() total_loss loss.item() scheduler.step() print(fEpoch {epoch1}, Loss: {total_loss:.4f})这里有一个细节我从train_dataset.data里直接取图片然后手动做归一化。这样做是为了和预先生成的教师 logits 的索引对齐。如果你用DataLoader重新加载索引顺序可能和保存 logits 的时候不一致导致软标签和图片对不上。这个坑很隐蔽一旦出现学生模型的准确率会莫名其妙地低而且很难排查。4.3 超参数搜索与结果对比为了找到最佳的超参数组合我做了一个简单的网格搜索。温度 T 取 [1, 2, 3, 4]alpha 取 [0.3, 0.5, 0.7]学生模型参数量取 [5000, 10000, 20000]。每组组合跑三次取平均准确率。总共跑了 4×3×3×3 108 次实验每次大概 30 秒总共花了不到一个小时。最终的实验结果表明最佳组合是 T2、alpha0.5、学生参数量 10000测试准确率 98.7%。这个结果比不用蒸馏的基线97.2%高了 1.5 个百分点比教师模型99.4%低了 0.7 个百分点。考虑到学生模型的参数量只有教师模型的 1/20这个压缩效果是可以接受的。配置Talpha学生参数量测试准确率基线无蒸馏--1000097.2%配置 A10.51000098.3%配置 B20.51000098.7%配置 C40.51000098.3%配置 D20.31000098.5%配置 E20.71000098.4%配置 F20.5500097.9%配置 G20.52000098.8%5. 常见问题与排查技巧实录5.1 学生模型准确率不升反降怎么办这是蒸馏里最常见的问题。如果你发现学生模型蒸馏后的准确率比不用蒸馏还低按照以下顺序排查第一检查教师模型的精度。如果教师模型的测试准确率低于 99%先把它训好再说。教师模型不行后面的一切都是白搭。第二检查软标签和图片的索引是否对齐。这个问题很隐蔽但后果很严重。你可以在训练前随机抽几张图片手动对比一下教师 logits 的 argmax 和真实标签是否一致。如果不一致说明索引乱了。第三检查温度系数和 alpha 权重。T 太大或者 alpha 太大都会导致硬标签的监督信号被削弱。先把 T 设成 1、alpha 设成 0.5看看效果然后再逐步调。第四检查学生模型的容量。如果学生模型太小它根本学不动软标签。先把学生模型放大一点确认蒸馏有效果再逐步压缩。5.2 torchvision 下载 MNIST 报 404 的解决方法这个问题我遇到过两次一次是在公司内网一次是在家里的宽带环境。表现是downloadTrue的时候直接抛异常提示 404 或者连接超时。解决办法有两个方法一手动下载。MNIST 的四个文件可以从多个公开镜像获取下载后放到./data/MNIST/raw/目录下确保文件名和 torchvision 期望的一致。然后把download参数改成False。方法二换用其他数据加载方式。比如用sklearn.datasets.fetch_openml(mnist_784)或者直接从本地 CSV 文件加载。不过这些方式得到的数据格式和 torchvision 的不一样需要额外做转换。提示如果你在多个环境里跑实验建议把 MNIST 数据放在一个固定的共享目录里所有实验都从这个目录加载避免重复下载和路径不一致的问题。5.3 蒸馏训练 loss 震荡的排查思路蒸馏训练的 loss 震荡比普通训练更常见因为涉及两个损失项的加权。如果 loss 震荡得厉害可以尝试以下方法降低学习率。蒸馏训练的学习率通常要比普通训练小一点因为软标签提供的梯度信号比较“柔和”学习率太大会导致参数更新过猛。增大 batch size。batch size 太小的话每个 batch 的软标签分布波动较大loss 自然不稳定。我一般用 128 或 256。检查温度系数。T 太小的话软标签的分布很尖锐KL 散度的梯度会很大T 太大的话分布太平滑梯度又太小。T2 到 T4 之间通常比较稳定。使用梯度裁剪。在loss.backward()之后、optimizer.step()之前加一行torch.nn.utils.clip_grad_norm_(student.parameters(), max_norm1.0)可以有效抑制梯度爆炸导致的 loss 震荡。5.4 常见问题速查表问题现象可能原因排查方法解决方案学生准确率低于基线教师模型没训好检查教师测试准确率重新训练教师模型至 99.2% 以上学生准确率低于基线软标签索引错位对比 logits argmax 和标签确保索引对齐学生准确率低于基线学生容量太小逐步增大参数量学生参数量至少为教师的 1/20loss 震荡严重学习率太大观察 loss 曲线降低学习率至 0.001 或更小loss 震荡严重batch size 太小检查 batch size增大至 128 或 256收敛速度慢温度系数太大打印软标签分布降低 T 至 2 左右收敛速度慢alpha 太大调整 alpha降低 alpha 至 0.5 左右下载 MNIST 报 404下载源不可用检查网络和 URL手动下载并放到 raw 目录6. 实操心得与避坑建议6.1 教师模型的选择比想象中重要很多人做蒸馏实验的时候随便拿一个模型当教师觉得“反正教师比学生强就行”。但实际做下来教师模型的质量对蒸馏效果的影响非常大。一个训练充分的教师模型输出的软标签分布是“有信息量”的——它知道哪些类别容易混淆哪些类别几乎不可能混淆。而一个训练不充分的教师模型输出的软标签分布是“糊”的学生学到的也是糊的东西。我的建议是教师模型的测试准确率至少要比学生模型的目标准确率高 1 个百分点以上而且教师模型的训练 loss 要基本收敛。在 MNIST 上教师模型的测试准确率最好在 99.3% 以上。6.2 温度系数和 alpha 需要联合调T 和 alpha 不是独立的它们之间有耦合关系。T 越大软标签分布越平滑KL 散度的梯度越小这时候可能需要适当增大 alpha 来补偿。反过来T 越小软标签分布越尖锐KL 散度的梯度越大alpha 可以适当减小。我一般会先固定 alpha0.5调 T找到最佳的 T 之后再固定 T调 alpha。这样比同时调两个参数要高效得多。在 MNIST 上T2、alpha0.5 是一个比较稳妥的起点。6.3 学生模型的容量不能压得太狠蒸馏的目的是压缩模型但压缩是有极限的。如果学生模型的容量太小它连基本的分类边界都拟合不好更别说学习软标签里的“暗知识”了。我的经验是学生模型的参数量至少要是教师模型的 1/20蒸馏才能稳定带来正向收益。如果压缩率超过 1/50蒸馏的效果会大打折扣甚至不如直接训练。当然这个比例不是绝对的和任务的复杂度有关。MNIST 比较简单1/20 就够了如果是 ImageNet 这种复杂任务可能需要 1/10 甚至更高。6.4 预生成软标签比在线蒸馏更高效在线蒸馏是指每个 batch 都跑一遍教师模型来生成软标签离线蒸馏是指提前把所有训练样本的软标签生成好存到磁盘上。我强烈建议用离线蒸馏原因有两个一是速度快教师模型只需要跑一遍后续训练学生模型的时候直接读软标签就行二是稳定软标签是固定的不会因为教师模型的 batch normalization 状态变化而波动。离线蒸馏的唯一缺点是占磁盘空间。MNIST 训练集有 6 万张图片每张图片的软标签是 10 个浮点数总共也就 2.4 MB完全可以接受。如果是更大的数据集可以考虑用半精度浮点数存储或者只存 top-k 个类别的概率。6.5 蒸馏不是万能的最后说一个我踩过的认知坑我一开始以为蒸馏能“无中生有”把一个小模型的准确率提升到接近教师模型的水平。但实际上蒸馏的增益是有限的。学生模型的准确率上限受限于它的容量蒸馏只是帮助学生更高效地利用训练数据而不是突破容量瓶颈。在 MNIST 上教师模型 99.4%学生模型蒸馏后 98.7%差了 0.7 个百分点。这 0.7 个百分点就是容量差距带来的“硬损失”蒸馏无法弥补。如果你需要更高的准确率要么增大教师模型要么增大学生模型没有别的捷径。这个内容后续还可以这样扩展把同样的蒸馏流程迁移到 CIFAR-10 上看看温度系数和 alpha 的最佳值会不会变化或者尝试不同的蒸馏损失函数比如 MSE 损失、注意力蒸馏、特征蒸馏对比它们在小模型上的效果差异。MNIST 只是一个起点真正的挑战在更复杂的数据集和更极端的压缩率上。