BERT知识蒸馏实战：用PyTorch实现轻量级NLP模型部署

1次阅读

共计 2072 个字符，预计需要花费 6 分钟才能阅读完成。

在自然语言处理（NLP）领域，BERT 等大型预训练模型虽然效果显著，但在实际工业部署中常面临两大难题：

内存占用高：BERT-base 模型参数达 110MB，难以在移动设备或嵌入式系统中运行
推理延迟大：单次文本分类推理耗时可达 100ms 以上，无法满足实时性要求高的场景

传统解决方案如模型剪枝（Pruning）和量化（Quantization）虽然能压缩模型体积，但往往伴随着明显的精度损失（通常下降 5 -15%）。知识蒸馏（Knowledge Distillation）通过迁移教师模型（Teacher Model）的 ” 知识 ” 到小型学生模型（Student Model），能在保持精度的同时显著减小模型体积。

方法	参数量压缩比	精度损失	硬件兼容性	实现复杂度
剪枝（Pruning）	60-70%	中	高	中
量化（Quantization）	75%	中高	中	低
知识蒸馏（Distillation）	80%+	低	高	高

graph TD
    A[教师模型 BERT-base] -->| 输出 logits 和隐藏层 | B[蒸馏损失函数]
    C[学生模型 BiLSTM] --> B
    B --> D[联合优化]

import torch
import torch.nn as nn
from transformers import BertModel

# 教师模型加载
teacher = BertModel.from_pretrained('bert-base-uncased')

# 学生模型定义
class StudentModel(nn.Module):
    def __init__(self, vocab_size, hidden_size=128):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size, bidirectional=True)
        self.fc = nn.Linear(hidden_size*2, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        x, _ = self.lstm(x)
        return self.fc(x[:, -1, :])

# 蒸馏损失函数
class DistillLoss(nn.Module):
    def __init__(self, temp=5.0, alpha=0.7):
        super().__init__()
        self.temp = temp
        self.alpha = alpha  # logits 损失权重
        self.kl_loss = nn.KLDivLoss(reduction='batchmean')
        self.ce_loss = nn.CrossEntropyLoss()

    def forward(self, student_logits, teacher_logits, labels):
        # logits 蒸馏
        soft_teacher = nn.functional.softmax(teacher_logits/self.temp, dim=-1)
        soft_student = nn.functional.log_softmax(student_logits/self.temp, dim=-1)
        kld_loss = self.kl_loss(soft_student, soft_teacher) * (self.temp**2)

        # 常规分类损失
        ce_loss = self.ce_loss(student_logits, labels)

        # 加权求和
        return self.alpha * kld_loss + (1-self.alpha) * ce_loss