从零掌握模型蒸馏技术：服务器上实战知识蒸馏

模型蒸馏技术应用越来越广泛，它能让一个复杂的大模型（老师）把自己的知识教给一个小模型（学生），从而实现模型压缩和推理加速。
本文面向零基础用户，以经典的 BERT 模型为例，一步步教你在服务器上完成一次完整的蒸馏训练。

准备工作：搭建蒸馏环境

在开始蒸馏之前，确保服务器已安装以下依赖（推荐使用 Python 3.8 以上版本）：

# 安装 PyTorch（根据 CUDA 版本选择）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装 Transformers 和 Datasets
pip install transformers datasets

# 安装训练常用库
pip install accelerate tensorboard

检查 GPU 是否可用：

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

如果返回 True，说明环境就绪。
如果你是纯 CPU 环境，也可以继续操作，但蒸馏速度会慢很多。

实战操作：BERT 知识蒸馏

我们使用 Hugging Face 的 distilbert-base-uncased 作为学生模型，bert-base-uncased 作为老师模型。
蒸馏的核心是让学生模型学习老师模型的输出分布（logits）以及真实标签。

1. 加载数据和模型

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset

# 老师模型
teacher = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
teacher.eval()

# 学生模型
student = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)

# 分词器
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 加载示例数据集（例如 IMDB 情感分类）
dataset = load_dataset("imdb", split="train[:1000]")
def tokenize(batch):
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=128)
dataset = dataset.map(tokenize, batched=True)

2. 定义蒸馏损失

蒸馏损失通常由两部分组成：学生与老师的 soft target 之间的 KL 散度，以及学生的预测与真实标签之间的交叉熵损失。

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    soft_targets = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction="batchmean"
    ) * (T * T)
    ce_loss = F.cross_entropy(student_logits, labels)
    return alpha * soft_targets + (1 - alpha) * ce_loss

3. 训练循环

使用 PyTorch 的标准训练流程：

from torch.utils.data import DataLoader
from transformers import AdamW

loader = DataLoader(dataset, batch_size=16, shuffle=True)
optimizer = AdamW(student.parameters(), lr=5e-5)

student.train()
for epoch in range(3):
    for batch in loader:
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["label"]

        with torch.no_grad():
            teacher_outputs = teacher(input_ids, attention_mask=attention_mask)
        student_outputs = student(input_ids, attention_mask=attention_mask)

        loss = distillation_loss(student_outputs.logits, teacher_outputs.logits, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    print(f"Epoch {epoch+1} loss: {loss.item():.4f}")

# 保存学生模型
student.save_pretrained("./student-distilled")

避坑指南：常见问题与应对

1. 显存不足：如果训练时 GPU 内存溢出，可减小 batch_size 或使用 gradient_accumulation_steps。另外，确保将老师模型设为 eval 模式并用 torch.no_grad() 包围。
2. Loss 不下降：检查学习率是否合适（建议 2e-5 到 5e-5），或者蒸馏温度 T 是否过大（一般设为 2-8）。
3. 加载模型报错：Hugging Face 的模型会自动下载到缓存，如果网络差，可先运行 transformers-cli download bert-base-uncased 手动下载。
4. 训练时间太长：可以先在小数据集上跑通流程，再扩展到全量数据。

效果验证：对比模型大小与推理速度

蒸馏完成后，比较两个模型：

# 查看参数量
torchinfo pytorch \"bert-base-uncased\" \"distilbert-base-uncased\"
# 或者直接在代码中计算

编写简单的推理脚本：

import time
sample = tokenizer("This movie is great!", return_tensors="pt")

# 老师推理
start = time.time()
teacher(**sample)
print(f"Teacher inference time: {time.time()-start:.4f}s")

# 学生推理
start = time.time()
student(**sample)
print(f"Student inference time: {time.time()-start:.4f}s")

通常学生模型参数量减少约 40%，推理速度提升 50% 以上，精度下降不到 2%。
这就是模型蒸馏技术应用的直接收益。

高频问题解答

Q：蒸馏后的模型怎么部署？
A：学生模型结构和普通 Hugging Face 模型一样，可以直接用 pipeline 加载或导出为 ONNX 进行生产部署。

Q：我可以蒸馏其他模型吗？
A：可以。蒸馏适用于任意 transformer 架构，只需修改 AutoModelForSequenceClassification 为对应的任务模型即可。

Q：没有 GPU 能蒸馏吗？
A：能，但训练会非常慢，建议使用 Google Colab 或租用云 GPU 实例。

如果你正在应用模型蒸馏技术，建议先按本文步骤完整执行，再根据自己的场景调整学生模型、温度和损失权重；
遇到异常时优先回看避坑部分。