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

1次阅读

没有评论

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

在自然语言处理（NLP）领域，BERT 等大型预训练模型凭借强大的语义理解能力成为主流选择。但当我们将这些模型部署到生产环境时，往往会遇到两个致命问题：

内存占用高 ：BERT-base 模型参数量达 1.1 亿，加载需要超过 1.2GB 内存
推理延迟大 ：单次推理在 CPU 上可能需要 500ms 以上，难以满足实时性要求

常见的模型压缩方案各有优劣：

Fine-tuning（微调）：仅调整最后一层，无法减少模型体积
Pruning（剪枝）：可能破坏模型结构完整性，需要复杂重训练
Distillation（蒸馏）：通过知识迁移训练小模型，兼顾尺寸与精度

教师模型 ：标准的 bert-base-uncased（12 层 Transformer）
学生模型 ：3 层 Transformer 结构，每层隐藏维度缩减为 512

import torch
from transformers import BertModel

# 定义蒸馏损失函数
def kl_div_loss(student_logits, teacher_logits, T=2.0):
    """
    KL 散度蒸馏损失实现
    Args:
        T: 温度系数，软化概率分布
    """
    soft_teacher = torch.nn.functional.softmax(teacher_logits/T, dim=-1)
    soft_student = torch.nn.functional.log_softmax(student_logits/T, dim=-1)
    return torch.nn.functional.kl_div(soft_student, soft_teacher, reduction='batchmean') * (T**2)

# 注意力矩阵蒸馏
class AttentionDistillLoss(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mse_loss = torch.nn.MSELoss()

    def forward(self, student_attn, teacher_attn):
        """
        student_attn: [batch, heads, seq_len, seq_len]
        teacher_attn: 同上
        """
        return self.mse_loss(student_attn, teacher_attn.detach())

# 初始化模型
teacher = BertModel.from_pretrained('bert-base-uncased')
student = TinyBert(config)

# 优化器设置
optimizer = torch.optim.AdamW(student.parameters(), lr=5e-5)
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=1000, num_training_steps=10000)

# 训练循环
for batch in dataloader:
    with torch.no_grad():
        teacher_outputs = teacher(**batch)

    student_outputs = student(**batch)

    # 计算三大损失
    cls_loss = kl_div_loss(student_outputs.logits, teacher_outputs.logits)
    attn_loss = attention_loss(student_outputs.attentions, teacher_outputs.attentions)
    hidden_loss = mse_loss(student_outputs.hidden_states, teacher_outputs.hidden_states)

    total_loss = 0.5*cls_loss + 0.3*attn_loss + 0.2*hidden_loss

    # 梯度裁剪
    torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
    optimizer.step()
    scheduler.step()