如何用知识蒸馏与 ONNX 压缩加速文本模型

BERT-base 级别的文本模型在 CPU 上单条推理几十毫秒，QPS 一上来机器成本直线上涨。压缩加速三板斧：蒸馏把知识搬进小模型、ONNX 把计算图优化掉、量化把位宽压下来——三者叠加常见 5～10 倍提速而指标只掉零点几个点。本文按顺序走通并给出回归方法。

    graph LR
    A[Teacher: BERT-base] -->|soft logits| B[Student: 4层小模型]
    B --> C[导出 ONNX 计算图]
    C --> D[INT8 动态量化]
    D --> E{延迟与指标回归}
    E -- 达标 --> F[替换线上推理]
    E -- 不达标 --> G[回退 FP32 ONNX 或加大 Student]

前置准备

以文本分类任务为例，准备已训好的 teacher 与无标签蒸馏数据：

pip install transformers torch onnxruntime datasets

蒸馏数据：线上真实请求日志脱敏 5～10 万条（比训练集分布更贴近服役分布）
Student 结构：与 teacher 同族但层数减半（如 4 层 hidden 384）

步骤一：蒸馏训练

student 同时拟合硬标签与 teacher 的 soft logits（暗知识在 soft 分布里）：

import torch.nn.functional as F

T = 4.0  # 温度:越高 soft 分布信息越平滑
for batch in distill_loader:
    with torch.no_grad():
        t_logits = teacher(**batch).logits
    s_logits = student(**batch).logits
    loss_kd = F.kl_div(F.log_softmax(s_logits / T, dim=-1),
                       F.softmax(t_logits / T, dim=-1), reduction="batchmean") * T * T
    loss_ce = F.cross_entropy(s_logits, batch["labels"])
    (0.7 * loss_kd + 0.3 * loss_ce).backward()
    opt.step(); opt.zero_grad()

先在小验证集确认 student 硬标签基线，再开始蒸馏
温度 T 从 2～4 试起；teacher 越强、数据越贴近线上，蒸馏收益越大

步骤二：导出 ONNX

把 student 导出为静态计算图，交给 ONNX Runtime 做算子融合优化：

python -m transformers.onnx --model=./student_distilled \
  --feature=sequence-classification onnx/student.onnx

import onnxruntime as ort
sess = ort.InferenceSession("onnx/student.onnx", providers=["CPUExecutionProvider"])
out = sess.run(None, {"input_ids": ids, "attention_mask": mask})

导出后用同一批输入对比 PyTorch 与 ORT 的输出 logits，最大绝对差应 < 1e-4。

步骤三：INT8 动态量化

CPU 部署场景下动态量化收益最大，一行配置：

from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic("onnx/student.onnx", "onnx/student_int8.onnx",
                 weight_type=QuantType.QUInt8)

动态量化：权重 INT8、激活运行时量化，无需校准集
GPU 部署则考虑 ORT 的 TensorRT EP 或静态量化（需校准数据）

步骤四：延迟与指标双回归

建一张对比表，四个配置同机同批数据测量：

import time
def latency(sess_or_model, batches):
    t0 = time.time()
    for b in batches: predict(sess_or_model, b)
    return round((time.time()-t0) / len(batches) * 1000, 2)  # ms/条