知识蒸馏实战：用PyTorch把BERT中文文本分类能力压缩进轻量模型-育师

简介：面向中文文本分类与模型压缩场景，这是一份基于Pytorch的知识蒸馏项目资源，适合希望掌握BERT蒸馏到轻量级BiLSTM模型的算法工程师和高年级学生。项目以Hugging Face上的bert-base-chinese为教师模型，将其在THUCNews十类中文语料上输出的logits作为软标签，蒸馏至BiLSTM学生模型，同时顺带加入梯度累加、混合精度训练、对抗训练等扩展实验，便于读者理解多种训练策略如何与蒸馏任务搭配。压缩包共43个文件，包括22个Python脚本、9个pkl模型/数据文件、5个txt说明、4个json配置、1个shell脚本及标签文件，整体约63.85MB。代码按data、config、models、checkpoints、processor目录组织，config内可控制训练、验证、测试、预测并选择是否启用上述扩展策略；processor分别处理BERT与BiLSTM两种输入格式，BiLSTM侧使用单字输入和5000字词表，结构清楚。目前已有312人学习，适合作为NLP分类与模型蒸馏方向的中高阶实践参考。

1. 知识蒸馏落地中文文本分类：把 BERT 的精度压进轻量模型

一条典型的落地场景：中文新闻分类，BERT 微调后验证集准确率 94%，但线上推理延迟 300ms，QPS 一起来就扛不住；换轻量模型直接训练，准确率掉到 89%。知识蒸馏的解法是让 BERT 当教师、轻量模型当学生，学教师的“判断轨迹”而不是照背标注。这套实践跑在 Pytorch 上，落在中文文本分类场景。

反直觉的点在于：蒸馏温度 T 不是越高越好，设到 10 以上学生反而什么都学不到；alpha 调错，精度甚至不如直接训练。这些参数背后有明确的计算逻辑，不是玄学。方案适合正在做人工智能大作业、需要完整可复现流程的人，也适合推理延迟吃紧、想压模型体积的工程师。人工智能正从尝鲜工具变日常帮手，“大模型带小模型”正是落地时最常用的一招。

2. 知识蒸馏的教师-学生框架：软标签、温度 T 与损失函数怎么选

2.1 为什么中文文本分类需要蒸馏：从 BERT 到轻量模型的收益

中文文本分类的任务形态通常比较简单：输入一段中文文本，输出一个类别标签。用 BERT 做效果好，是因为它在预训练阶段见过大量中文语料，对字词语义关系有先验理解；但 BERT 参数量动辄上亿，单条前向在 CPU 上要几百毫秒，GPU 上也要几十毫秒。生产环境不会只看准确率，推理成本同样参与选型，这就逼着工程师在精度和延迟之间做取舍。

知识蒸馏的核心思路是教师-学生框架：先训练一个高精度的教师模型，再训练一个结构更简单的学生模型，让学生模型的输出不断逼近教师模型的输出。学生不再直接学 one-hot 标注，而是学教师给出的概率分布。这个分布里含有类别间“像不像”的信息——一段关于篮球的文本，教师可能给出体育 0.7、娱乐 0.2，学生学到这层关系后，对边界样本的判别会比直接训练更稳。

这个方案在中文文本分类里很常见，原因有二。一是中文分类数据集往往不大，标注成本高，直接训练小模型容易过拟合，蒸馏等于把大模型从海量语料里学到的先验迁移给小模型；二是蒸馏对数据链路改动小，教师和学生模型用同一套 tokenizer 和 Dataset，预处理代码不用重写，评估脚本可以复用到部署环节。

2.2 软标签与温度 T：蒸馏损失的三个关键参数

第一个要理解的是软标签。教师对每个样本输出的 logits 经 softmax 后得到概率分布，这个分布就是软标签。与硬标签相比，软标签携带类别相似度信息，这是学生模型能超过直接训练的关键。但直接拿教师 softmax 输出当标签有个问题：教师模型很有信心时，正确类别概率接近 1，其他类别接近 0，软标签退化成硬标签，信息量丢失。

解决办法是引入温度 T：先对 logits 除以 T 再做 softmax。T 大于 1 时概率分布变平滑，教师对错误类别的“宽容度”被放大；T 等于 1 就是普通 softmax；T 趋近 0 时分布变成 one-hot。实际项目中 T 一般取 2 到 8，我习惯先固定 T=4 跑通，再按验证集精度微调。与其听别人说“T 越大越好”，不如自己扫一遍曲线。

第三个参数是软标签损失在总损失里的权重 alpha。蒸馏总损失由软标签的 KL 散度和硬标签的交叉熵加权构成，公式是：L = alpha * T^2 * KL(softmax(student_logits / T), softmax(teacher_logits / T)) + (1 - alpha) * CE(student_logits, hard_label)。KL 散度的参数顺序容易写反：Pytorch 的 KLDivLoss 要求 input 是 log-probabilities、target 是 probabilities，所以代码里要用 log_softmax(student_logits / T) 作为 input，用 softmax(teacher_logits / T) 作为 target，这个顺序错了 loss 数值会乱但没有语法报错，是典型的黑匣子问题。

注意：T^2 不是拍脑袋乘的。对 student_logits 求梯度时，log_softmax(student_logits / T) 的梯度会缩小 1/T 倍，交叉熵的梯度还会再引入一次 1/T，整体梯度缩了 T^2 倍，需要乘回去补偿，否则 T 越大梯度越小，学生根本学不动。

2.3 教师模型选型：微调过的 BERT 与预训练权重的差别

教师模型必须是对当前任务微调过的。蒸馏迁移的是“任务相关知识”，不是通用语言知识。没微调的 BERT 对文本分类输出几乎全是均匀分布，软标签里没有判别信息，学生学不到东西。常见做法是先让教师模型在训练集上微调到验证集精度不再上升，然后冻结参数，再开始蒸馏。

教师精度也不是越高越好。教师完全过拟合训练集时，它在训练样本上的软标签信心过高，类别间相似度信息反而被压缩。判断办法很简单：看教师训练集和验证集的精度差距，差距超过 5 个百分点就说明过拟合明显，先给教师加早停或正则，再回来调蒸馏参数。这个点经常被忽略，大部分人遇到学生精度不升，第一反应是调 alpha 和 T，其实问题出在教师身上。

3. 用 Pytorch 搭出可复现的蒸馏流水线：环境、数据与教师模型

3.1 Pytorch 环境搭建：conda 创建虚拟环境的最小命令

整个项目建议单独建虚拟环境，别把 torch 装进系统 Python。常见做法是 conda 建 python 3.9 环境，再走 pip 安装。下面这组命令是一套干净的最小安装流程：

conda create -n distill python=3.9 -y conda activate distill pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install transformers datasets scikit-learn

第一行创建名为 distill 的虚拟环境，第二行激活。pip 加 --index-url 是指定从 Pytorch 官方 wheel 源安装，cu118 表示 CUDA 11.8 版；纯 CPU 环境去掉这个参数直接 pip install torch 就行。transformers 负责加载 BERT 与 tokenizer，datasets 用来读取常见中文分类数据集，scikit-learn 用来算评估指标。装完后建议顺手验证一遍：

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

Windows 上 pip 安装 Pytorch 最经典的报错是 conda activate 提示“无法将 conda 项识别为 cmdlet、函数、脚本文件”，多数原因是 Anaconda 的 Scripts 目录没进 PATH，或者没跑过 conda init。我一般直接执行 conda init powershell 再重开终端，比手动改 PATH 可靠。验证命令输出里 torch.version是预期版本、cuda.is_available() 是 True，就说明环境 OK。这里多花两分钟，后面能省一晚上的排错时间。

3.2 中文文本分类数据准备：字级别 tokenization 与 label 映射

中文文本分类第一个决策是 tokenization 粒度。常见做法有两种：字符级别，每个汉字一个 token，适合 BiLSTM 这类轻量学生模型；或者直接用 BERT 自带的 tokenizer，WordPiece 对中文会切成字或子词，不需要额外分词。中文数据没有英文那种空格分词的习惯，拿 jieba 先分词再喂 BERT 反而是多余操作，还会引入分词错误。

import json import torch from torch.utils.data import Dataset class TextClassificationDataset(Dataset): def __init__(self, file_path, tokenizer, max_len=128): self.tokenizer = tokenizer self.max_len = max_len self.samples = [] with open(file_path, "r", encoding="utf-8") as f: for line in f: obj = json.loads(line) self.samples.append((obj["text"], obj["label"])) def __len__(self): return len(self.samples) def __getitem__(self, idx): text, label = self.samples[idx] encoded = self.tokenizer( text, max_length=self.max_len, padding="max_length", truncation=True, return_tensors="pt", ) return { "input_ids": encoded["input_ids"].squeeze(0), "attention_mask": encoded["attention_mask"].squeeze(0), "label": torch.tensor(label, dtype=torch.long), }

这个 Dataset 把每行 JSON 中的 text 和 label 读出来，经 tokenizer 编码成固定长度 128。padding 到 max_length 是为了让批量训练时 shape 一致；truncation 截断超长文本。中文新闻标题大多在几十个字以内，128 足够；长文档分类需要提到 256 或 512。label 必须是 0 开始的连续整数。如果原始数据里 label 是中文类别名，先用一个固定字典做映射，比如{'体育': 0, '财经': 1, '娱乐': 2}，映射一旦确定就不要改，训练、蒸馏、评估三个阶段共用同一份。

3.3 训练教师模型：以 BERT 为教师的最小训练脚本

教师模型用 transformers 的 BertForSequenceClassification 加载 bert-base-chinese。bert-base-chinese 的词表本身就是字级别，中文数据可以直接把原始文本扔给 tokenizer，不用额外分词。最小训练脚本如下：

from transformers import BertForSequenceClassification, BertTokenizer, AdamW tokenizer = BertTokenizer.from_pretrained("bert-base-chinese") teacher = BertForSequenceClassification.from_pretrained( "bert-base-chinese", num_labels=num_classes ) teacher.to(device) teacher.train() optimizer = AdamW(teacher.parameters(), lr=2e-5) for epoch in range(3): for batch in train_loader: optimizer.zero_grad() input_ids = batch["input_ids"].to(device) attn_mask = batch["attention_mask"].to(device) labels = batch["label"].to(device) outputs = teacher(input_ids, attention_mask=attn_mask, labels=labels) outputs.loss.backward() optimizer.step()

loss 由 transformers 内部算出，是标准的 softmax 交叉熵。用 AdamW 而不是 Adam，因为 BERT 的权重衰减策略需要 AdamW 内置的 decoupled weight decay。学习率 2e-5 是 BERT 微调的经验区间 1e-5 到 5e-5 里的常见取值。num_labels 必须等于类别数，否则分类头维度对不上；训练轮次 3 是保底，数据量小或类别多时可以跑到 5 轮，用验证集精度做早停判断。

教师训完先保存，后面蒸馏阶段直接加载，不用重训。

teacher.save_pretrained("./teacher_model") tokenizer.save_pretrained("./teacher_model")

保存教师模型时把 tokenizer 一起存。加载的时候一行代码还原：BertForSequenceClassification.from_pretrained('./teacher_model')。后面蒸馏循环里教师模型要固定参数，这一点在第 4 章展开。

4. 蒸馏训练学生模型：软标签损失与硬标签损失的联合优化

4.1 学生模型结构选择：BiLSTM 还是小 BERT

学生模型的结构决定蒸馏结果的上限。我常用的两条路线是双向 LSTM 和蒸馏版小 BERT。双向 LSTM 参数量在几百万量级，CPU 上单条推理 1~3ms，训练一个 epoch 只要几分钟；小 BERT（bert-tiny、bert-mini）保留了 Transformer 结构，和教师同源，知识迁移更顺滑，但体积和延迟都比 LSTM 大。结构选型本质是精度与延迟的权衡。

这些数值是经验区间，不是固定结论。中文文本分类数据集小的时候，BiLSTM 蒸馏后反而可能比小 BERT 更稳，因为小 BERT 结构复杂，数据量不够时容易过拟合。我一般先跑 BiLSTM 作为基线，如果学生精度离教师差距超过 3 个点再换 bert-tiny。下面给一个 BiLSTM 学生模型的紧凑定义：

import torch import torch.nn as nn class BiLSTMClassifier(nn.Module): def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes, num_layers=2): super().__init__() self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0) self.lstm = nn.LSTM( embed_dim, hidden_dim, num_layers=num_layers, batch_first=True, bidirectional=True ) self.classifier = nn.Linear(hidden_dim * 2, num_classes) def forward(self, input_ids, attention_mask=None): x = self.embedding(input_ids) x, _ = self.lstm(x) x = x[:, -1, :] return self.classifier(x)

embedding 的 padding_idx=0 要和 tokenizer 的 pad_token_id 对应，这样 padding 位置的向量在反向传播时梯度为 0，不参与更新；LSTM 用双向，所以 classifier 输入维度是 hidden_dim * 2，取序列最后一个时间步的输出。attention_mask 在这里没有真正参与计算，如果要做严格 mask，需要把 padding 位置的 hidden state 归零再取 pooling，工程上为了速度许多人直接忽略这一点，效果差别不大。vocab_size 直接用 tokenizer.vocab_size。

4.2 蒸馏主循环：温度 T、alpha 与教师 logits 的冻结

蒸馏主循环是整个项目最核心的代码。关键约束是：教师模型必须 eval 模式且不更新梯度，否则教师参数被学生梯度污染，两个模型一起乱跑，精度必跌。下面这段可以直接跑：

import torch import torch.nn.functional as F T = 4.0 alpha = 0.7 teacher_model.eval() student_model.train() for batch in train_loader: input_ids = batch["input_ids"].to(device) attn_mask = batch["attention_mask"].to(device) labels = batch["label"].to(device) with torch.no_grad(): teacher_logits = teacher_model(input_ids, attention_mask=attn_mask).logits student_logits = student_model(input_ids, attn_mask) loss_soft = F.kl_div( F.log_softmax(student_logits / T, dim=-1), F.softmax(teacher_logits / T, dim=-1), reduction="batchmean", ) * (T * T) loss_hard = F.cross_entropy(student_logits, labels) loss = alpha * loss_soft + (1 - alpha) * loss_hard optimizer.zero_grad() loss.backward() optimizer.step()

教师输出在 no_grad 上下文里只做前向；soft target 是 teacher_logits 除以 T 后 softmax，student_logits 除以 T 后 log_softmax 作为 KL 散度的 input；最后乘 T*T 补偿梯度。hard loss 用普通交叉熵，保证学生不偏离真实标注。

三个参数值得单独强调。温度 T 必须同时作用在两个模型上，只除一边会导致两个分布在完全不同的尺度上做 KL 散度，loss 数值很小但方向是错的，学生学到的分布会被严重拉偏。reduction 用 batchmean 而不是 mean，mean 是逐元素求平均再除以 batch 外的维度数，batchmean 才是对整个 batch 求均值，配合乘 T*T 时梯度量纲才正确。alpha=0.7 表示 70% 的损失来自软标签；太高学生完全模仿教师、教师有错它也学错，太低蒸馏退化成直接训练。

注意：KLDivLoss 的 input 和 target 顺序不能反。input 是学生 logits 的 log_softmax，target 是教师 logits 的 softmax。写反不会报错，但 loss 数值会变成负数甚至 NaN，排查起来相当痛苦。

4.3 评估与模型保存：验证集精度与模型体积对比

蒸馏训练完要做两件事：验证集精度对比和模型体积对比。精度对比证明蒸馏有效，体积和延迟对比说明部署收益。先跑评估：

def evaluate(model, loader): model.eval() correct = 0 total = 0 with torch.no_grad(): for batch in loader: input_ids = batch["input_ids"].to(device) attn_mask = batch["attention_mask"].to(device) labels = batch["label"].to(device) logits = model(input_ids, attention_mask=attn_mask) preds = logits.argmax(dim=-1) correct += (preds == labels).sum().item() total += labels.size(0) return correct / total teacher_acc = evaluate(teacher_model, val_loader) student_acc = evaluate(student_model, val_loader) print(f"teacher_acc={teacher_acc:.4f}, student_acc={student_acc:.4f}")

保存学生模型时，我会连同 label_map 和配置一起存，而不是只存 state_dict：

torch.save({ "model_state_dict": student_model.state_dict(), "label_map": label_map, "config": {"num_classes": num_classes, "max_len": max_len}, }, "student_distilled.pt")

保存时把 label_map 和 max_len 一起带上，加载模型做推理时才能还原 tokenizer 参数和标签映射，不然导出 ONNX 后根本不知道输出节点对应哪个类别。这是工程化落地的习惯——实验里很多人只存 state_dict，上线时才发现缺少 metadata，还得重新对齐训练时的配置。模型体积对比也很简单：分别打印教师和学生模型的 checkpoint 文件大小，再跑一次推理 benchmark 记录延迟。如果学生精度比教师低不到 2 个点，推理快了 10 倍以上，这个蒸馏就是值得的。

5. 中文文本分类知识蒸馏的避坑指南：从 NaN 到精度不升的 5 个排查点

蒸馏项目跑通不难，难的是精度上不去之后不知道从哪里排查。下面这 5 个坑我都在中文文本分类场景里踩过，按现象、原因、解决的顺序写，方便你对号入座。

5.1 现象：loss 出现 NaN 或剧烈震荡

原因最常见的是学生模型学习率过大。学生模型随机初始化，用教师模型 2e-5 的学习率偏小，但调到 1e-3 又可能过头；其次是 KL 散度里没加 log_softmax，直接把概率传进去，log 0 得 -inf。解决：学生模型学习率从 1e-3 起步，观察训练 loss，出现 NaN 就降 10 倍；同时检查代码里 KL 散度的 input 是不是 log_softmax 的输出。还可以给 loss 加 clamp：loss_soft = loss_soft.clamp(max=10.0)，防止单条异常样本的 loss 炸掉整个 batch。这个 clamp 不影响正常训练，但能让你从 NaN 的恐慌里先解脱出来。

5.2 现象：蒸馏后学生精度反而低于直接训练

先别怀疑蒸馏方案有问题，大概率是参数没配对。alpha=0.9 时学生几乎不学硬标签，教师在校验集上犯错的位置学生全部继承；T 太小软标签接近硬标签，蒸馏意义消失。另一个原因是教师过拟合训练集，它的软标签在训练样本上信心过高，类别间相似度信息被压缩。解决：alpha 降到 0.5~0.6，T 升到 6 再扫一遍；教师训练时加早停，别把验证集精度练到顶点再停。很多时候学生精度不升是教师学过了头，不是学生的问题。这个结论我用两个数据集验证过，几乎每次都成立。

5.3 现象：教师 logits 与学生 logits 维度不匹配

transformers 模型输出的是 ModelOutput 对象，拿 outputs.logits 才是 logits，拿 hidden_states 或 pooler_output 维度会变成 (batch, hidden_size)，和学生的 (batch, num_classes) 对不上。另一个隐蔽点：BertForSequenceClassification 的默认 num_labels=2，忘了传 num_classes，教师分类头维度就是 2，跟学生维度不匹配。解决：在拼接蒸馏损失前打印两个 logits 的 shape，确认都是 (batch, num_classes)。这个检查 10 秒能做完，能省掉后面所有维度的报错。

5.4 现象：中文文本编码不一致导致标签错位

这个坑最隐蔽，乱码不一定报错。训练文件是 GBK 编码，Windows 上默认编码读出来是乱码，tokenizer 把乱码切成无意义的字，学生模型等于在学噪声，loss 照样下降但精度上不去。另一个场景是 label 映射没统一，训练时 label 是字符串，蒸馏时转 int 映射出错，标签错位但 loss 看起来正常。解决：打开文件时显式写 encoding='utf-8'，统一数据文件编码；label_map 用一个全局字典定义一次，训练、验证、蒸馏三处引用同一个对象，禁止临时改映射。这个习惯是从一次评估精度忽高忽低的排查里换来的。

5.5 现象：GPU 显存不足与 batch size 的选择

蒸馏阶段教师和学生两个模型同时在前向，显存占用接近教师单独训练的两倍；BERT 的 attention 是输入长度的平方复杂度，max_len 从 128 提到 256，显存占用直接翻倍。解决：蒸馏时给教师模型单独用小 batch，学生 64、教师 32，教师 forward 完结果就释放；max_len 不要盲目调大，中文短文本 128 够用。更实用的做法是提前把教师 logits 离线算好存成 npy 文件，蒸馏循环里只加载 logits，完全避开教师模型占显存的问题，也不用再维护 no_grad 的逻辑。数据集不大的时候这个方案我强烈推荐。

6. 进阶验证：温度扫描、教师退化与模型导出

6.1 温度扫描与教师熵检查

模型跑通后，第一件值得做的验证是温度扫描：固定 alpha=0.7，T 从 2 到 8 每个值完整训练一个 epoch，记录学生验证集精度画曲线。我一般拿 T=1 当基线，它等于无蒸馏直接训练；如果 T=4 时精度还不如 T=1，优先查教师模型过拟合而不是继续加 T。扫描前顺手算一下教师 logits 的 softmax 熵：如果绝大多数样本熵低于 0.1，说明教师太自信、软标签信息量低，这时把 T 提到 8 或 10 反而有帮助。这个判断比盲目调参快得多。

6.2 学生模型导出：从 Pytorch 到 ONNX

学生模型训练完成、评估通过后，就可以导出做线上部署。我一般用 torch.onnx.export 转 ONNX，然后交给 onnxruntime 推理：

torch.onnx.export( student_model, (dummy_input_ids, dummy_attn_mask), "student.onnx", input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes={ "input_ids": {0: "batch"}, "attention_mask": {0: "batch"}, }, opset_version=11, )

dynamic_axes 声明 batch 维度是可变的，线上推理时单个请求和合并批量都兼容。ONNX 导出后在 CPU 上通常比 Pytorch eager 模式快 20%~50%，配合 onnxruntime 的 intra_op_num_threads 参数还能进一步压延迟。导出前确认学生模型已经调成 eval 模式，torch.onnx.export 不支持包含 dropout 的 train 状态。这一步做完，学生模型才真正具备上线的形态，前面所有蒸馏工作才算落地。

我最早做蒸馏时习惯把 T 和 alpha 一起调，每次都两个参数一块改，验证集精度上去了就以为是蒸馏有效，换了一个数据集就翻车。后来养成固定一个扫另一个的习惯，才把蒸馏从黑匣子变成可预期的流程。这套基于 Pytorch 的中文文本分类蒸馏流程，核心就是三件事：训好教师、选对学生、把损失函数的三个参数调对。希望帮到你。

本文还有配套的精品资源，点击获取