知识蒸馏实战：从零构建MNIST分类器并对比不同损失函数效果

1. 知识蒸馏入门从MNIST手写数字识别开始知识蒸馏听起来高大上但其实原理很简单——就像老师教学生一样。想象一下你有一个经验丰富的老师模型通常比较大、准确率高现在要训练一个轻量级的学生模型。通过知识蒸馏学生模型不仅能学习标准答案hard target还能学习老师对各类别的思考方式soft target。MNIST手写数字数据集是入门的最佳选择。这个数据集包含0-9的手写数字图片每张图片都是28x28的灰度图。我们先用一个复杂的MLP网络训练教师模型再用一个简单的MLP网络作为学生模型。你会发现通过知识蒸馏训练的学生模型比直接训练的学生模型表现更好。为什么选择MLP网络因为它结构简单训练速度快非常适合教学演示。教师网络我们设计为3层全连接1200-1200-10学生网络则是精简版20-20-10。虽然MNIST很简单但这个实验流程可以迁移到更复杂的数据集和模型上。2. 环境搭建与数据准备2.1 安装必要的Python库首先确保你安装了最新版的PyTorch。我推荐使用conda创建虚拟环境conda create -n distillation python3.8 conda activate distillation pip install torch torchvision tqdm torchinfo这些库的作用分别是torch: PyTorch深度学习框架torchvision: 提供MNIST等标准数据集tqdm: 显示训练进度条torchinfo: 打印模型结构信息2.2 数据加载模块实现数据预处理是机器学习的第一步。对于MNIST我们只需要简单的ToTensor转换import torchvision from torchvision import transforms from torch.utils.data import DataLoader def load_data(batch_size128): transform transforms.Compose([ transforms.ToTensor(), ]) train_set torchvision.datasets.MNIST( root./data, trainTrue, downloadTrue, transformtransform ) test_set torchvision.datasets.MNIST( root./data, trainFalse, downloadTrue, transformtransform ) train_loader DataLoader(train_set, batch_sizebatch_size, shuffleTrue) test_loader DataLoader(test_set, batch_sizebatch_size, shuffleFalse) return train_loader, test_loader这里有几个实用技巧batch_size设置为128能平衡内存占用和训练稳定性训练集需要shuffle测试集则不需要数据会自动下载到./data目录3. 模型设计与实现3.1 教师模型构建教师模型要足够强大这里我们设计一个三层的MLPimport torch.nn as nn class TeacherModel(nn.Module): def __init__(self, input_dim784, hidden_dim1200, output_dim10): super().__init__() self.fc1 nn.Linear(input_dim, hidden_dim) self.fc2 nn.Linear(hidden_dim, hidden_dim) self.fc3 nn.Linear(hidden_dim, output_dim) self.dropout nn.Dropout(0.5) self.relu nn.ReLU() def forward(self, x): x x.view(x.size(0), -1) # 展平输入 x self.relu(self.fc1(x)) x self.dropout(x) x self.relu(self.fc2(x)) x self.dropout(x) x self.fc3(x) return x关键设计点使用了Dropout(0.5)防止过拟合隐藏层维度设为1200确保模型容量足够ReLU激活函数提供非线性能力3.2 学生模型设计学生模型要小巧很多class StudentModel(nn.Module): def __init__(self, input_dim784, hidden_dim20, output_dim10): super().__init__() self.fc1 nn.Linear(input_dim, hidden_dim) self.fc2 nn.Linear(hidden_dim, hidden_dim) self.fc3 nn.Linear(hidden_dim, output_dim) self.relu nn.ReLU() def forward(self, x): x x.view(x.size(0), -1) x self.relu(self.fc1(x)) x self.relu(self.fc2(x)) x self.fc3(x) return x学生模型特点隐藏层只有20个神经元是教师的1/60没有Dropout层因为模型本身已经很简单同样使用ReLU激活函数4. 训练流程实现4.1 常规训练方法我们先实现标准的训练流程作为baselineimport torch import torch.optim as optim from tqdm import tqdm def train_model(model, train_loader, test_loader, epochs50, lr0.0001): device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lrlr) best_acc 0.0 for epoch in range(epochs): model.train() train_loss 0.0 for inputs, labels in tqdm(train_loader): inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() # 测试集评估 model.eval() correct 0 total 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() acc correct / total print(fEpoch {epoch1}/{epochs}, Loss: {train_loss/len(train_loader):.4f}, Acc: {acc:.4f}) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_model.pth) print(fBest Accuracy: {best_acc:.4f}) return model4.2 知识蒸馏训练这才是重头戏。知识蒸馏的关键在于损失函数设计def distill_train(teacher, student, train_loader, test_loader, epochs50, lr0.0001, temp7.0, alpha0.3): device torch.device(cuda if torch.cuda.is_available() else cpu) teacher teacher.to(device).eval() # 教师模型设为评估模式 student student.to(device) hard_loss nn.CrossEntropyLoss() soft_loss nn.KLDivLoss(reductionbatchmean) optimizer optim.Adam(student.parameters(), lrlr) best_acc 0.0 for epoch in range(epochs): student.train() train_loss 0.0 for inputs, labels in tqdm(train_loader): inputs, labels inputs.to(device), labels.to(device) with torch.no_grad(): teacher_logits teacher(inputs) student_logits student(inputs) # 计算hard loss student_hard_loss hard_loss(student_logits, labels) # 计算soft loss (知识蒸馏损失) soft_student F.log_softmax(student_logits/temp, dim1) soft_teacher F.softmax(teacher_logits/temp, dim1) distillation_loss soft_loss(soft_student, soft_teacher) # 组合损失 loss alpha * student_hard_loss (1-alpha) * temp * temp * distillation_loss optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() # 测试集评估 student.eval() correct 0 total 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels inputs.to(device), labels.to(device) outputs student(inputs) _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() acc correct / total print(fEpoch {epoch1}/{epochs}, Loss: {train_loss/len(train_loader):.4f}, Acc: {acc:.4f}) if acc best_acc: best_acc acc torch.save(student.state_dict(), best_student.pth) print(fBest Accuracy: {best_acc:.4f}) return student关键参数说明temp(温度)控制预测分布的平滑程度通常设为3-10alphahard loss和soft loss的权重平衡KL散度损失衡量学生和教师输出分布的差异5. 不同损失函数效果对比5.1 ChatGPT版实现这是我们前面使用的标准实现效果最好soft_student F.log_softmax(student_logits/temp, dim1) soft_teacher F.softmax(teacher_logits/temp, dim1) distillation_loss F.kl_div(soft_student, soft_teacher, reductionbatchmean) loss alpha * hard_loss (1-alpha) * temp * temp * distillation_loss特点学生输出用log_softmax教师输出用softmax乘以temp平方平衡梯度大小5.2 同济子豪兄版实现distillation_loss F.kl_div( F.softmax(student_logits/temp, dim1), F.softmax(teacher_logits/temp, dim1), reductionbatchmean ) loss alpha * hard_loss (1-alpha) * temp * temp * distillation_loss问题学生输出也用了softmax可能导致数值不稳定实际测试中loss有时会出现负数5.3 文心一言版实现def distillation_loss(student_logits, teacher_logits, temperature): student_probs F.softmax(student_logits / temperature, dim1) teacher_probs F.softmax(teacher_logits / temperature, dim1) kl_divergence F.kl_div( student_probs.log(), teacher_probs, reductionbatchmean ) * (temperature ** 2) return kl_divergence * temperature特点额外乘以了temperature可能使损失值过大hard loss和蒸馏损失不在一个量级5.4 实验结果对比方法最佳准确率训练稳定性收敛速度普通训练93.83%高快ChatGPT版95.86%高中等同济子豪兄版92.87%中(有负loss)慢文心一言版92.90%高慢从实验结果看ChatGPT版的实现效果最好能让学生模型达到接近教师模型的准确率教师模型最佳为98.69%。文心一言版虽然稳定但效果提升有限。同济子豪兄版则存在数值不稳定的问题。6. 实战技巧与调优建议6.1 温度参数(temp)的选择温度参数是知识蒸馏的核心超参数低温度(1-3)强调困难样本的学习中等温度(5-10)平衡软硬目标高温度(10)过度平滑失去有用信息建议从temp3开始尝试逐步增加到7-10。可以通过观察教师输出的分布来选择# 查看教师输出的分布 with torch.no_grad(): logits teacher(inputs) probs F.softmax(logits/temp, dim1) print(probs)理想的分布应该正确类别的概率明显高于其他但其他类别之间也有区分度6.2 alpha参数的调整alpha控制hard loss和soft loss的权重alpha1退化为普通训练alpha0完全依赖教师知识推荐0.1-0.5之间我的经验是早期训练alpha0.3-0.5后期微调alpha0.1-0.26.3 训练策略优化两阶段训练先用较高的temp和较低的alpha训练后期降低temp提高alpha学习率调整scheduler optim.lr_scheduler.StepLR(optimizer, step_size20, gamma0.1)早停机制if acc best_acc: best_acc acc patience 0 else: patience 1 if patience 10: break7. 进阶应用与扩展7.1 跨架构知识蒸馏我们的例子是MLP→MLP但知识蒸馏可以应用于CNN→MLPResNet→MobileNetBERT→DistilBERT关键是要确保学生和教师的输出维度一致适当调整温度参数7.2 使用中间层特征除了输出层的知识还可以蒸馏中间层特征# 在模型中定义特征提取hook def get_features(name): def hook(model, input, output): features[name] output.detach() return hook teacher.fc2.register_forward_hook(get_features(fc2)) student.fc2.register_forward_hook(get_features(fc2)) # 在损失函数中加入特征蒸馏损失 feature_loss F.mse_loss(student_features[fc2], teacher_features[fc2])7.3 现成工具库推荐MMClassification: 提供多种知识蒸馏算法RepDistiller: 包含12种最新蒸馏方法HuggingFace Transformers: 支持BERT等模型的蒸馏例如使用RepDistiller:from distiller import DistillKL, HintLoss criterion_kd DistillKL(temp4) criterion_hint HintLoss() # 计算损失 loss_kd criterion_kd(student_output, teacher_output) loss_hint criterion_hint(student_feature, teacher_feature) loss alpha * hard_loss beta * loss_kd gamma * loss_hint8. 完整代码实现与实验日志8.1 完整训练脚本import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader from tqdm import tqdm import numpy as np import random # 设置随机种子 def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True set_seed() # 数据加载 def load_data(batch_size128): transform transforms.Compose([ transforms.ToTensor(), ]) train_set datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) test_set datasets.MNIST(./data, trainFalse, transformtransform) train_loader DataLoader(train_set, batch_sizebatch_size, shuffleTrue) test_loader DataLoader(test_set, batch_sizebatch_size, shuffleFalse) return train_loader, test_loader # 模型定义 class TeacherModel(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(784, 1200) self.fc2 nn.Linear(1200, 1200) self.fc3 nn.Linear(1200, 10) self.dropout nn.Dropout(0.5) self.relu nn.ReLU() def forward(self, x): x x.view(x.size(0), -1) x self.relu(self.fc1(x)) x self.dropout(x) x self.relu(self.fc2(x)) x self.dropout(x) x self.fc3(x) return x class StudentModel(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(784, 20) self.fc2 nn.Linear(20, 20) self.fc3 nn.Linear(20, 10) self.relu nn.ReLU() def forward(self, x): x x.view(x.size(0), -1) x self.relu(self.fc1(x)) x self.relu(self.fc2(x)) x self.fc3(x) return x # 知识蒸馏训练 def distill_train(teacher, student, train_loader, test_loader, epochs50, lr0.0001, temp7.0, alpha0.3): device torch.device(cuda if torch.cuda.is_available() else cpu) teacher teacher.to(device).eval() student student.to(device) hard_loss nn.CrossEntropyLoss() soft_loss nn.KLDivLoss(reductionbatchmean) optimizer optim.Adam(student.parameters(), lrlr) best_acc 0.0 for epoch in range(epochs): student.train() train_loss 0.0 for inputs, labels in tqdm(train_loader, descfEpoch {epoch1}/{epochs}): inputs, labels inputs.to(device), labels.to(device) with torch.no_grad(): teacher_logits teacher(inputs) student_logits student(inputs) # 计算损失 student_hard_loss hard_loss(student_logits, labels) soft_student F.log_softmax(student_logits/temp, dim1) soft_teacher F.softmax(teacher_logits/temp, dim1) distillation_loss soft_loss(soft_student, soft_teacher) loss alpha * student_hard_loss (1-alpha) * temp * temp * distillation_loss optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() # 测试集评估 student.eval() correct 0 total 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels inputs.to(device), labels.to(device) outputs student(inputs) _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() acc correct / total print(fLoss: {train_loss/len(train_loader):.4f}, Acc: {acc:.4f}) if acc best_acc: best_acc acc torch.save(student.state_dict(), best_student.pth) print(fBest Accuracy: {best_acc:.4f}) # 主程序 if __name__ __main__: # 加载数据 train_loader, test_loader load_data() # 训练教师模型 print(Training Teacher Model...) teacher TeacherModel() train_model(teacher, train_loader, test_loader) # 知识蒸馏训练学生模型 print(\nDistilling Student Model...) student StudentModel() distill_train(teacher, student, train_loader, test_loader)8.2 典型实验日志Training Teacher Model... Epoch 1/50: 100%|██████████| 469/469 [00:0400:00, 100.23it/s] Loss: 0.3245, Acc: 0.9453 ... Epoch 50/50: 100%|██████████| 469/469 [00:0400:00, 101.12it/s] Loss: 0.0124, Acc: 0.9869 Best Accuracy: 0.9869 Distilling Student Model... Epoch 1/50: 100%|██████████| 469/469 [00:0500:00, 85.12it/s] Loss: 2.4573, Acc: 0.9021 ... Epoch 50/50: 100%|██████████| 469/469 [00:0500:00, 86.34it/s] Loss: 0.8745, Acc: 0.9586 Best Accuracy: 0.9586从日志可以看出教师模型最终准确率达到98.69%通过知识蒸馏小模型也能达到95.86%的准确率相比直接训练的93.83%知识蒸馏带来了2%的提升9. 常见问题与解决方案9.1 损失函数出现NaN值可能原因温度参数设置过小学习率过高梯度爆炸解决方案# 梯度裁剪 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) # 检查温度 assert temp 1.0, Temperature should be greater than 1.0 # 使用更稳定的损失计算 soft_teacher F.softmax(teacher_logits/temp 1e-8, dim1) # 添加小常数9.2 学生模型表现不如教师可能原因模型容量差距过大蒸馏温度不合适alpha参数设置不当调试步骤逐步增加学生模型的隐藏层维度尝试不同的温度值(3,5,7,10)调整alpha从0.1到0.59.3 训练速度慢优化建议使用更大的batch size(256或512)启用混合精度训练scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): student_logits student(inputs) # 计算损失... scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()使用DataLoader的num_workers参数DataLoader(..., num_workers4, pin_memoryTrue)10. 实际应用中的经验分享在工业级应用中我发现几个实用技巧渐进式蒸馏先用大温度训练逐步降低温度temps [10, 7, 5, 3] # 分阶段训练 for temp in temps: distill_train(..., temptemp)多教师蒸馏结合多个教师模型的知识teacher1_logits teacher1(inputs) teacher2_logits teacher2(inputs) soft_teacher (F.softmax(teacher1_logits/temp, dim1) F.softmax(teacher2_logits/temp, dim1)) / 2注意力蒸馏不仅蒸馏输出还蒸馏注意力图# 在模型中定义注意力层 class AttentionLayer(nn.Module): def __init__(self): super().__init__() self.attn nn.Linear(hidden_dim, 1) def forward(self, x): attn_weights F.softmax(self.attn(x), dim1) return attn_weights # 在损失函数中加入注意力损失 attn_loss F.mse_loss(student_attn, teacher_attn)在线蒸馏教师和学生模型同步更新# 教师模型也参与训练 teacher.train() student.train() teacher_logits teacher(inputs) student_logits student(inputs) # 教师模型用真实标签训练 teacher_loss hard_loss(teacher_logits, labels) optimizer_teacher.zero_grad() teacher_loss.backward() optimizer_teacher.step() # 学生模型用蒸馏损失训练 ...这些技巧在我的多个项目中都有显著效果特别是在模型压缩和移动端部署场景下。知识蒸馏不仅能提升小模型性能还能发现教师模型中隐含的知识表示。