知识蒸馏实战教程：用PyTorch将大模型压缩为小模型

“什么时候，蒸馏我自己！”——看到这个标题，很多同学可能会会心一笑。这句话表面上是一句程序员的自我调侃，但拆开来看，它恰好指向了深度学习里一个非常实用的技术方向：知识蒸馏（Knowledge Distillation）。通俗点说，知识蒸馏就是“训练一个大模型当老师，再让一个小模型跟着老师学”，最终让小模型在参数量大幅减少的情况下，尽量逼近大模型的精度表现。本文会用一篇完整的实战教程，把知识蒸馏的原理、代码、训练过程和踩坑点讲清楚，让你看完之后不仅能理解“蒸馏”到底在蒸什么，还能自己动手跑通一个完整的蒸馏训练流程。

这篇文章适合以下几类读者：

• 刚入门深度学习，想了解模型压缩与加速的同学。
• 已经能跑通基础 PyTorch 训练，但还不清楚如何做“师生模型”训练的人。
• 在端侧、边缘设备或 Web 端部署模型时，发现模型太大、推理太慢，想通过蒸馏压缩模型的工程师。

读完本文，你将掌握知识蒸馏的核心思想、损失函数的构造方法、温度系数 T 的作用，以及一套基于 PyTorch 的完整可运行示例。下面我们正式开始。

1. 背景与核心概念

1.1 为什么需要知识蒸馏

先来看一个常见的现实问题：在图像分类、文本分类等任务中，通常模型越大、层数越深，效果越好。一个 ResNet-50、BERT-base 甚至更大的模型，在 GPU 服务器上训练和推理都没有问题。但一旦要部署到手机 App、浏览器、嵌入式设备上，就会遇到几个棘手的问题：

• 模型参数量太大，占用存储空间。
• 推理延迟高，用户点击一次要等好几秒。
• 内存占用过高，低端设备直接崩溃。

有人会说，那直接换一个小模型不就行了？确实，MobileNet、SqueezeNet 这类轻量网络在结构上做了很多优化，但如果直接用小型化数据去训练一个轻量模型，效果通常比大模型差一截。原因是小模型容量有限，很难从零开始完全学到复杂数据中的规律。

知识蒸馏提供了一种折中方案：先用大模型（教师网络）在数据上学到丰富的知识，然后让小模型（学生网络）去模仿大模型的输出。因为教师网络的输出包含了“类间相似度”这种软信息，学生网络能学到的内容，比只依赖真实标签（One-Hot 硬标签）要多得多。

1.2 什么是知识蒸馏

知识蒸馏这个概念最早由 Hinton 等人在 2015 年的论文Distilling the Knowledge in a Neural Network中正式提出。它的核心思路可以概括为三步：

1. 训练一个性能较好的大模型，称为教师网络（Teacher）。
2. 设计一个参数量更少的小模型，称为学生网络（Student）。
3. 在学生网络的训练过程中，不仅让它学习真实标签，还要让它模仿教师网络的输出概率分布。

这里的关键在于“概率分布”。普通的分类任务使用交叉熵损失，目标是把真实类别的概率推向 1，其他类别推向 0。但教师网络在预测时，除了正确类别之外，其他类别的概率并不是正好为 0，而是有高有低。比如一张图片是一只猫，教师网络可能输出：

这个概率分布里，“狗”和“老虎”的概率比“其他”高，说明教师网络认为猫和狗、猫和老虎在视觉特征上有一定的相似性。这种相似性就是教师网络从海量数据中学到的“软知识”。学生网络如果只学习硬标签，就只会得到一个“猫”的判断结果；但如果学习教师网络的软输出，它就能额外知道“猫有点像狗，但不太像老虎”。这种软知识，正是蒸馏能够提升小模型效果的根本原因。

1.3 知识蒸馏的常见应用场景

知识蒸馏并不是只能用于图像分类。在工程实践中，它的应用场景非常广泛：

• 模型压缩：将大规模 Transformer 或 ResNet 蒸馏为小模型，用于移动端部署。
• 跨架构迁移：把 Transformer 的知识蒸馏到 LSTM 或 CNN 结构中。
• 多任务与大模型简化：用一个大模型同时训练出多个专用小模型。
• 半监督与自监督增强：利用教师模型对无标注数据生成伪标签，结合蒸馏训练学生模型。
• 推荐系统与排序模型：用复杂排序模型蒸馏出轻量召回模型。

理解了背景之后，我们接下来从环境准备开始，一步步完成一个知识蒸馏的实战项目。

2. 环境准备与版本说明

2.1 基础环境

本文的示例代码基于 Python 和 PyTorch。建议使用以下环境：

• Python：3.8 或更高版本。
• PyTorch：1.10 及以上，2.x 版本也可以。
• torchvision：与 PyTorch 版本对应的版本。
• 操作系统：Windows / Linux / macOS 均可。

如果你使用的是 GPU 环境（CUDA 11.x 或更高版本），训练速度会快很多。如果没有 GPU，用 CPU 训练也能完成本文的示例，只是耗时会长一些。

创建虚拟环境并安装依赖：

# 创建虚拟环境（可选） python -m venv distill_env # 激活虚拟环境 # Windows: distill_env\Scripts\activate # Linux/macOS: source distill_env/bin/activate # 安装依赖 pip install torch torchvision matplotlib

版本说明：不同版本的 PyTorch 在 API 上略有差异，但本文用到的nn.CrossEntropyLoss、nn.KLDivLoss、F.log_softmax都是长期稳定的接口，在各版本中均可使用。如果你使用的 PyTorch 版本非常新，遇到个别 API 变动，以官方文档为准。

2.2 示例项目结构

为了方便阅读和运行，我们采用下面的项目结构：

knowledge_distillation/ ├── train.py # 完整训练脚本 ├── models.py # 教师网络与学生网络定义 ├── distill.py # 蒸馏训练逻辑 └── README.md # 项目说明

为了让代码更清晰，这里把模型定义、蒸馏训练逻辑和主脚本分开。实际项目中可以按照自己的习惯组织目录结构，但保持模块化是一个好习惯。

3. 知识蒸馏核心原理拆解

3.1 软标签与温度系数 T

前面提到，教师网络输出的概率分布就是“软标签”。但这里有一个问题：如果教师网络的输出过于自信，比如正确类别的概率是 0.99，其他类别几乎为 0，那么软标签和硬标签的区别就不明显了，学生网络也学不到太多额外信息。

为了解决这个问题，Hinton 在论文中引入了“温度系数 T”。在计算概率分布时，不再直接使用网络输出的 logits 做 Softmax，而是先除以 T：

[ q_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)} ]

其中 (z_i) 是网络输出的原始 logits，T 是温度系数。

当 T=1 时，这就是普通的 Softmax 输出；当 T>1 时，概率分布变得更“平滑”，类别之间的差异被缩小，模型就能暴露出更多“软知识”。下图是一个简单的示意：

• T=1：输出概率分布比较尖锐。
• T=4：输出概率分布更平均，类间关系更明显。
• T 越大，分布越接近均匀分布，但超过一定范围后，有效信息反而被淹没。

在蒸馏训练中，教师网络和学生网络在计算蒸馏损失时，通常使用相同的温度 T，且 T 往往大于 1（例如 T=4 或 T=5）。

3.2 蒸馏损失函数的构成

知识蒸馏的总损失由两部分组成：

1. 硬标签损失（Student Loss）：学生网络的输出与真实标签之间的交叉熵损失。
2. 蒸馏损失（Distill Loss）：学生网络的高温输出与教师网络的高温输出之间的 KL 散度损失。

数学表达可以写成：

[ L = \alpha \cdot L_{hard} + (1 - \alpha) \cdot L_{soft} ]

其中：

• (L_{hard}) 是学生网络输出与真实标签的交叉熵。
• (L_{soft}) 是学生网络与教师网络的软标签之间的 KL 散度。
• (\alpha) 是硬标签损失的权重，通常设置在 0.1 到 0.7 之间。

为什么需要两个损失？如果只看蒸馏损失，学生网络只学到“模仿老师”；如果只看硬标签损失，那就退化成普通训练。两者结合，学生既能从老师的经验中受益，又不会完全被老师的错误判断带偏。

3.3 为什么用 KL 散度而不是交叉熵

KL 散度（Kullback-Leibler Divergence）用于衡量两个概率分布之间的差异。在蒸馏中，教师网络的输出概率作为“目标分布”，学生网络的输出概率作为“预测分布”，KL 散度越小，说明学生越接近老师。

计算方式如下：

# 教师网络和学生网络的输出都先经过 log_softmax loss_kl = nn.KLDivLoss(reduction='batchmean')( F.log_softmax(student_output / T, dim=1), F.softmax(teacher_output / T, dim=1) ) * (T * T)

注意这里有一个细节：F.log_softmax在前，F.softmax在后。因为KLDivLoss的输入要求第一个参数是对数概率，第二个参数是普通概率。最后的(T * T)是梯度缩放修正，从论文中沿用下来的处理方式，目的是让损失值在不同温度下保持合理的尺度。

3.4 蒸馏训练的整体流程

一次完整的蒸馏训练过程可以概括为以下步骤：

1. 加载数据集，分成训练集和测试集。
2. 定义教师网络和学生网络。
3. 固定教师网络参数，先训练教师网络，或直接加载一个预训练好的教师模型。
4. 在每一轮训练中，同时向学生网络输入真实标签和教师网络的软输出。
5. 计算硬标签损失和蒸馏损失，加权求和后反向传播更新学生网络参数。
6. 在测试集上评估学生网络效果。

接下来就用 PyTorch 把上面这套流程完整实现出来。

4. 完整实战案例：PyTorch 实现知识蒸馏

4.1 创建项目结构

首先创建项目目录和文件：

mkdir knowledge_distillation cd knowledge_distillation touch train.py models.py distill.py

4.2 定义教师网络和学生网络

为了演示方便，我们使用 MNIST 手写数字数据集。MNIST 有 10 个类别，图像尺寸是 28x28，结构简单，训练速度快，非常适合用来验证蒸馏流程。

这里教师网络使用一个参数量较多的 CNN，学生网络使用一个参数量较少的 CNN。完整定义放在models.py中：

# 文件路径：models.py import torch import torch.nn as nn import torch.nn.functional as F class TeacherNet(nn.Module): """ 教师网络：参数量较多、表达能力更强。 结构：两层卷积 + 三层全连接 """ def __init__(self): super(TeacherNet, self).__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(64 * 7 * 7, 512) self.fc2 = nn.Linear(512, 256) self.fc3 = nn.Linear(256, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x class StudentNet(nn.Module): """ 学生网络：参数量少、结构轻量。 结构：一层卷积 + 两层全连接 """ def __init__(self): super(StudentNet, self).__init__() self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(16 * 14 * 14, 64) self.fc2 = nn.Linear(64, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.fc2(x) return x

说明：

• 教师网络采用两层卷积加三层全连接，参数约 50 万级别。
• 学生网络只保留一层卷积加两层全连接，参数约 3 万级别。
• 两个网络的输入输出维度一致，都是 MNIST 的 1x28x28 输入，10 类输出。

4.3 编写蒸馏训练逻辑

接下来在distill.py中编写蒸馏训练的核心逻辑。这里包括训练教师网络、蒸馏训练学生网络两个环节。因为教师网络定义较复杂，我们先训练它，蒸馏时冻结它的参数。

# 文件路径：distill.py import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms def train_teacher(model, device, train_loader, optimizer, criterion, epoch): """ 普通训练教师网络 """ model.train() running_loss = 0.0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() running_loss += loss.item() if batch_idx % 100 == 0: print(f'Epoch {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] ' f'Loss: {loss.item():.6f}') def evaluate(model, device, test_loader): """ 评估模型准确率 """ model.eval() correct = 0 total = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() total += target.size(0) acc = 100.0 * correct / total print(f'Test accuracy: {acc:.2f}%') return acc def distill_train(student, teacher, device, train_loader, optimizer, T, alpha): """ 蒸馏训练学生网络。 参数说明： - student: 学生网络 - teacher: 教师网络，训练完成后冻结 - T: 温度系数 - alpha: 硬标签损失的权重 """ student.train() teacher.eval() # 教师网络固定，不参与梯度更新 hard_criterion = nn.CrossEntropyLoss() soft_criterion = nn.KLDivLoss(reduction='batchmean') running_loss = 0.0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() # 学生网络输出 student_output = student(data) # 教师网络输出，不计算梯度 with torch.no_grad(): teacher_output = teacher(data) # 硬标签损失 loss_hard = hard_criterion(student_output, target) # 蒸馏损失：使用高温 softmax loss_soft = soft_criterion( F.log_softmax(student_output / T, dim=1), F.softmax(teacher_output / T, dim=1) ) * (T * T) # 总损失 loss = alpha * loss_hard + (1 - alpha) * loss_soft loss.backward() optimizer.step() running_loss += loss.item() if batch_idx % 100 == 0: print(f'Distill [Batch {batch_idx}] Total Loss: {loss.item():.6f}, ' f'Hard Loss: {loss_hard.item():.6f}, Soft Loss: {loss_soft.item():.6f}')

这段代码是蒸馏训练的核心，有几个地方值得仔细说明：

1. teacher.eval()确保教师网络中的 BatchNorm/Dropout 不产生影响。推理模式下必须设置。
2. with torch.no_grad()包裹教师网络的前向传播，节省显存和计算量。
3. KL 散度损失中，教师网络输出经过softmax，学生网络输出经过log_softmax，顺序不能颠倒。
4. 最终损失是硬标签损失和蒸馏损失的加权和，通过alpha控制两者的权重。

4.4 编写主训练脚本

主脚本train.py负责加载数据、初始化模型、依次训练教师网络和学生网络。这里使用 MNIST 数据集，如果本地没有数据，PyTorch 会自动下载。

# 文件路径：train.py import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from models import TeacherNet, StudentNet from distill import train_teacher, evaluate, distill_train def main(): # 设备设置 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'Using device: {device}') # 数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) # 加载 MNIST train_dataset = datasets.MNIST( root='./data', train=True, download=True, transform=transform ) test_dataset = datasets.MNIST( root='./data', train=False, download=True, transform=transform ) batch_size = 128 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) # 初始化模型 teacher = TeacherNet().to(device) student = StudentNet().to(device) # ========== 第一步：训练教师网络 ========== print('========== Training Teacher Network ==========') teacher_optimizer = optim.Adam(teacher.parameters(), lr=0.001) teacher_criterion = nn.CrossEntropyLoss() for epoch in range(5): train_teacher(teacher, device, train_loader, teacher_optimizer, teacher_criterion, epoch + 1) train_acc = evaluate(teacher, device, train_loader) test_acc = evaluate(teacher, device, test_loader) print(f'Teacher Epoch {epoch + 1}: train_acc={train_acc:.2f}%, test_acc={test_acc:.2f}%') # ========== 第二步：蒸馏训练学生网络 ========== print('========== Distilling Student Network ==========') student_optimizer = optim.Adam(student.parameters(), lr=0.001) T = 4 # 温度系数 alpha = 0.3 # 硬标签损失权重 # 冻结教师网络 for param in teacher.parameters(): param.requires_grad = False for epoch in range(5): distill_train(student, teacher, device, train_loader, student_optimizer, T, alpha) test_acc = evaluate(student, device, test_loader) print(f'Student Distill Epoch {epoch + 1}: test_acc={test_acc:.2f}%') # 保存模型 torch.save(student.state_dict(), 'student_model.pth') torch.save(teacher.state_dict(), 'teacher_model.pth') print('Models saved.') if __name__ == '__main__': main()

4.5 运行与验证

在项目目录下运行：

python train.py

预期输出大致如下：

Using device: cuda ========== Training Teacher Network ========== Epoch 1 [0/60000] Loss: 0.358264 ... Test accuracy: 98.21% Teacher Epoch 1: train_acc=97.52%, test_acc=98.21% ... ========== Distilling Student Network ========== Distill [Batch 0] Total Loss: 1.982103, Hard Loss: 1.823456, Soft Loss: 0.553211 ... Student Distill Epoch 5: test_acc=97.26%

为了验证蒸馏的有效性，可以再单独训练一个不使用蒸馏的 StudentNet，并对比测试准确率。通常来说，经过蒸馏的学生网络会比从零训练的学生网络有 0.5 到 2 个百分点的提升，具体取决于数据集、温度和损失权重的设置。

5. 常见问题与排查思路

5.1 教师网络输出过拟合

如果教师网络在训练集上效果很好，但在测试集上泛化很一般，那么学生学到的是教师“背下来的答案”，而不是“理解的规律”。这种情况常见于教师网络过度训练，或者训练轮次过多。

建议：教师网络训练到验证集准确率不再提升时就停止，并尽量使用早停或正则化手段。

5.2 温度 T 设置不当

温度 T 太小，软标签接近硬标签，蒸馏退化为普通训练。温度 T 太大，软标签接近均匀分布，学生网络学不到类别间差异。

推荐做法：将 T 作为超参数，在验证集上分别尝试 2、3、4、5，找到最合适的值。

5.3 KL 散度损失出现 NaN

如果 KL 散度损失出现 NaN，通常是因为教师网络或学生网络的输出经过 Softmax 后出现极端值，或者 logits 太小导致数值不稳定。

排查顺序：

1. 检查数据是否归一化。
2. 检查网络输出是否有限值，可以在前向传播后打印output.min()和output.max()。
3. 尝试在 softmax 前对 logits 做 clip，例如限制在 [-5, 5]。
4. 降低学习率。

5.4 蒸馏后学生网络没有提升

这是最常见的问题。可能的原因有：

• 教师网络不够强，软知识价值有限。
• 学生网络结构太简单，容量严重不足，无论怎么学都学不会。
• 温度 T 设置过低，软标签和硬标签几乎一样。
• alpha 权重过高，蒸馏损失的影响微乎其微。
• 训练轮次不足，学生网络还没收敛就提前结束。

排查时可以先做一个 base line：直接训练学生网络并记录准确率，然后对比蒸馏训练后的准确率。如果两者几乎一样，优先调整 T 和 alpha。

5.5 显存不足

蒸馏训练需要同时加载教师网络和学生网络，显存开销比单独训练学生网络要高。可以尝试：

• 减少 batch size。
• 使用 CPU 训练（MNIST 这类小数据集完全可行）。
• 分阶段推理：预先用教师网络生成所有样本的软标签，保存到磁盘，蒸馏训练时直接读取，不需要在内存中保留教师网络。这也是工程上推荐的优化方式。具体思路是训练教师网络后，遍历训练集，把教师输出的软标签存为.npy或.pt文件，训练学生时只加载这些文件。

6. 最佳实践与工程建议

知识蒸馏在学术与工程中的实现方式远比上面的示例丰富，这里整理一些实用的工程建议。

6.1 先做模型容量评估

蒸馏不是万能的。如果教师网络和学生网络的容量差距过大，学生网络可能根本无法吸收教师的知识。建议先用下面的思路评估：

• 单独训练学生网络，看它能达到的准确率上限。
• 用教师网络结构替换学生网络，确定目标效果的“天花板”。
• 如果两者差距非常大，考虑增加学生网络容量，或者改用更优的轻量结构，而不是一味地加大蒸馏损失权重。

6.2 软标签预处理与离线蒸馏

在大型数据集上训练时，如果每个 epoch 都要前向传播教师网络，开销非常大。工程上更常见的方式是：

• 训练好教师网络后，对全部训练数据执行一次前向传播。
• 将教师网络的 softmax(T) 输出保存到磁盘，作为预计算软标签。
• 训练学生网络时，不再加载教师网络，直接从磁盘读取软标签。

这样做的好处有两点：一是显著减少显存和训练耗时；二是软标签可以重复使用，更换温度 T 或损失权重时无需重新跑教师网络。

6.3 动态权重与课程蒸馏

比较进阶的做法是动态调整 alpha。训练初期，学生网络还很弱，可以适当提高蒸馏损失权重，让它多观察教师的行为；训练中后期，学生逐渐成熟，再慢慢提高硬标签损失权重，让它与真实任务对齐。

简单实现伪代码如下：

# 动态 alpha 示例 alpha = max(0.7 - epoch * 0.1, 0.1)

6.4 集成蒸馏与自蒸馏

除了单教师蒸馏，还有几种常见的扩展方式：

• 多教师蒸馏：多个教师网络同时提供软标签，取平均或加权融合，学生网络可以获得更全面的知识。
• 自蒸馏：让同一个网络的大版本来蒸馏小版本，例如在训练过程中使用 EMA（指数滑动平均）模型作为教师。
• 特征蒸馏：不只模仿输出概率，还让学生网络模仿教师网络中间层的特征图，常用于检测、分割等任务。

这些方案都是建立在基础蒸馏框架之上的优化。先跑通本文的基础代码，再逐步引入复杂度，是比较稳妥的学习路径。

6.5 验证集与超参管理

蒸馏训练有两个网络、多个超参数，很容易过拟合验证集。建议建立一套规范：

• 固定教师网络，只调学生网络的超参数。
• 使用独立的验证集调参，测试集只在最终评估时使用一次。
• 记录每次实验的温度 T、权重 alpha、学习率、教师准确率、学生准确率，方便对比。

7. 总结与学习路线

本文从“什么时候，蒸馏我自己”这句调侃出发，系统梳理了知识蒸馏的核心概念、原理与完整实战。相信你现在已经能回答这几个关键问题：

• 蒸馏到底蒸的是什么：教师网络软输出中包含的类间相似性信息。
• 温度 T 的作用：控制概率分布的平滑程度，暴露更多软知识。
• 总损失如何构成：硬标签的交叉熵损失 + 教师软标签的 KL 散度损失。
• 如何用 PyTorch 实现：训练教师网络后冻结，再训练学生网络。

接下来，如果你想继续深入这个方向，可以从三个维度扩展：

1. 模型结构角度：尝试将教师学生模型替换为不同架构，观察蒸馏效果的差异。
2. 任务角度：将示例从 MNIST 换成 CIFAR-10，或者尝试在 NLP 文本分类任务中使用蒸馏。
3. 部署角度：将训练好的学生模型转换为 ONNX 或 TorchScript，在端侧推理框架中部署。

知识蒸馏并不是高不可攀的前沿技术，它更像是一套“站在巨人肩膀上学习”的训练范式。希望这篇文章能帮你跑通第一条蒸馏训练流水线，并在实际项目中真正用起来。如果本文对你有帮助，欢迎收藏备用。你也可以在自己的数据集上调整温度 T 和权重 alpha，看看学生网络能在多大程度上接近教师网络的表现。

动手试一下，跑出你自己的蒸馏模型，下一次再问“什么时候蒸馏我自己”的时候，你就能自信地回答：现在就可以。