BERT模型压缩实战：从理论到轻量化部署

1次阅读

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

BERT-base 模型拥有 1.1 亿参数（110M），在推理时需要占用约 1.2GB 显存。这样的资源消耗在以下场景中会带来显著挑战：

移动端应用：手机内存有限，大型模型难以直接部署
边缘设备：IoT 设备计算能力弱，需要轻量级模型
实时服务：高并发场景下，大模型推理延迟和成本激增

方法	压缩率	精度损失	计算开销	适用场景
8-bit 量化	4x	<3%	低	通用硬件部署
4-bit 量化	8x	5-8%	中	极度资源受限环境
结构化剪枝	2-5x	3-10%	高	需要保持矩阵运算
非结构化剪枝	5-10x	5-15%	高	高稀疏性需求
Logits 蒸馏	2-4x	2-6%	中	有高质量教师模型
中间层蒸馏	2-4x	1-5%	高	需要保留中间特征

from transformers import BertModel, BertTokenizer
import torch

# 加载原始模型
model = BertModel.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# 量化准备
def calibrate(model, calib_data):
    model.eval()
    with torch.no_grad():
        for batch in calib_data:
            inputs = tokenizer(batch, return_tensors='pt', padding=True)
            model(**inputs)

# 执行量化
quantized_model = torch.quantization.quantize_dynamic(
    model, 
    {torch.nn.Linear}, 
    dtype=torch.float16
)

from transformers import BertForSequenceClassification
import torch.nn.utils.prune as prune

model = BertForSequenceClassification.from_pretrained('bert-base-uncased')

# 自动选择阈值
parameters_to_prune = [(module, 'weight') 
    for module in model.modules() 
    if isinstance(module, torch.nn.Linear)
]

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.4  # 剪枝 40% 权重
)

import torch.nn as nn
import torch.nn.functional as F

class DistillLoss(nn.Module):
    def __init__(self, temp=5.0):
        super().__init__()
        self.temp = temp

    def forward(self, student_logits, teacher_logits):
        student_probs = F.log_softmax(student_logits/self.temp, dim=-1)
        teacher_probs = F.softmax(teacher_logits/self.temp, dim=-1)
        return F.kl_div(student_probs, teacher_probs, reduction='batchmean') * (self.temp**2)

在 GLUE 的 MRPC 任务上测试结果：