BERT模型压缩-量化实战

ai大模型-探索者2026-09-19 10:30

BERT 模型量化实战教程：把 409MB 的大模型压到 100MB

配套项目：今日头条新闻 10 分类（BERT 微调模型）。

一、为什么需要模型压缩

你辛苦训练的 BERT 模型有 409MB，效果很好（准确率 95%+），但是------

问题 影响
模型太大 服务器内存装不下、加载慢
推理太慢 用户等几百毫秒才出结果
成本高 服务器、电费、云费用都贵

结论：训练出的"又大又准"的模型，往往不是最终上线用的模型。 工业界上线前的最后一环，就是模型压缩------在尽量不损失精度的前提下，让模型变小变快。

模型压缩三大方法

方法 类比 原理 效果
量化 高清图→压缩图 降低数字精度（32位→8位） 体积÷4
蒸馏 导师教学生 大模型教小模型 体积÷几十
剪枝 修剪果树 删掉冗余参数 体积变小

本教程专注讲第一种：量化。

二、什么是量化（Quantization）

2.1 直观理解

想象一张高清照片 （4MB）和一张压缩照片（1MB）：

  • 高清：每个像素用很多位存储颜色，细节丰富
  • 压缩：精度降低，但肉眼看几乎没差别

量化就是干这个 ：模型的权重原本用 32 位浮点数 存储（精度极高、占 4 字节），量化后变成 8 位整数（精度降低、只占 1 字节）。

2.2 数字上发生了什么

复制代码 
量化前（FP32，32位浮点）:
  0.7838291927...    ← 精确，但占 4 字节

量化后（INT8，8位整数）:
  0.78               ← 粗略，但只占 1 字节

2.3 为什么精度几乎不掉？

因为神经网络非常"抗造"------它本来就有大量冗余参数。把 0.7838291927 变成 0.78，对最终结果的影响微乎其微。

一句话：量化 = 把"高清大图"换成"压缩小图"，人眼看不出差别，但体积小了 4 倍。

2.4 量化的两种类型

类型 说明 适用
动态量化（Dynamic） 推理时动态把权重量化，代码最简单 LSTM / Transformer 等
静态量化（Static） 训练后校准，精度更好但需要标定数据 CNN 等

本教程使用 动态量化 （torch.quantization.quantize_dynamic），一行搞定，最适合入门。

三、环境准备

bash 复制代码 
pip install torch transformers

确认环境：

python 复制代码 
import torch
print(torch.__version__)
print(torch.cuda.is_available())   # 量化在 CPU 上做即可

四、项目结构

复制代码 
06-model-compression/
└── bert_quantization/
    ├── bert_classifer_model.py      # BERT 分类模型定义
    ├── config.py                    # 配置（路径、超参）
    ├── utils.py                     # 数据加载
    ├── train.py                     # 训练 + 评估函数
    ├── bert_model_quantization.py   # ★ 量化主程序
    ├── predict_fun.py               # 量化模型的预测函数
    └── models_save/
        ├── bert20250521_.pt                    # 原始模型（输入）
        └── quantized_model_xxx.pt              # 量化模型（输出）

五、量化前置条件

量化不是重新训练 ，而是对已训练好的模型做压缩。所以你必须先有：

  1. 训练好的 BERT 分类模型 （如 bert_classifier_model.pt）
  2. 模型定义代码 （bert_classifer_model.py，量化后还要用它加载）
  3. 测试数据（用来对比量化前后精度）

六、模型定义（bert_classifer_model.py）

python 复制代码 
import torch
from transformers import BertModel
from config import Config

conf = Config()

class BertClassifier(torch.nn.Module):
    """BERT + 全连接分类层"""
    def __init__(self):
        super(BertClassifier, self).__init__()
        self.bert = BertModel.from_pretrained(conf.bert_path)      # BERT 主体
        self.fc = torch.nn.Linear(conf.hidden_size, conf.num_classes)  # 分类头 768→10

    def forward(self, input_ids, attention_mask):
        _, pooled = self.bert(input_ids=input_ids,
                              attention_mask=attention_mask,
                              return_dict=False)
        out = self.fc(pooled)     # 输出 10 个类别的分数
        return out

量化时重点处理的就是这里的 nn.Linear 全连接层（BERT 中参数量最大的部分）。

七、核心：量化代码（bert_model_quantization.py）

python 复制代码 
# bert_model_quantization.py
from bert_classifer_model import BertClassifier
from config import Config
import torch
from utils import build_dataloader
from train import model2dev

# 初始化配置
conf = Config()

if __name__ == '__main__':
    # 1. 创建数据迭代器（用于评估量化后精度）
    print('加载数据...')
    train_dataloader, test_dataloader, dev_dataloader = build_dataloader()

    # 2. 加载训练好的模型
    print("加载模型...")
    device = conf.device
    model = BertClassifier()
    model_path = conf.model_save_path
    # map_location='cpu' 指定映射设备；量化在 CPU 上进行
    model.load_state_dict(torch.load(model_path, map_location='cpu'))
    model.eval()

    print("查看量化前的模型结构=========================")
    print(model)

    # 3. ★ 核心：动态量化，把 nn.Linear 层转成 8 位整数
    quantized_model = torch.quantization.quantize_dynamic(
        model,
        {torch.nn.Linear},      # 量化哪些层类型
        dtype=torch.qint8       # 量化成 8 位整数
    )
    print("量化后的模型=========================")
    print(quantized_model)

    # 4. 评估量化后模型精度（可选，取消注释运行）
    # report, f1score, accuracy, precision = model2dev(quantized_model, test_dataloader, device)
    # print("Test Classification Report:", report)
    # print("Test F1:", f1score)
    # print("Test Accuracy:", accuracy)

    # 5. 计算 8-bit 量化后模型的内存占用（MB）
    quantized_size = sum(
        p.numel() * p.element_size() for p in quantized_model.parameters()
    ) / 1024 ** 2
    print(f"8-bit 量化后的模型内存: {quantized_size:.2f} MB")

    # 6. 保存整个量化模型
    torch.save(quantized_model, conf.quantized_model_save_path)
    print("保存量化模型成功！地址为：", conf.quantized_model_save_path)

八、代码逐行深度解析

① torch.load(model_path, map_location='cpu')

加载训练好的模型权重。

  • map_location='cpu'：把权重映射到 CPU（量化通常在 CPU 上做，避免设备不匹配）

② ★ 核心量化函数

python 复制代码 
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)
参数 含义
model 要量化的模型
{torch.nn.Linear} 量化哪些层------这里指定全连接层（BERT 参数量最大的部分）
dtype=torch.qint8 量化数据类型：8 位整数（原 FP32 是 32 位）

这一行就是量化的全部核心------PyTorch 帮你自动完成了所有参数的类型转换。

③ 计算内存占用

python 复制代码 
sum(p.numel() * p.element_size() for p in quantized_model.parameters()) / 1024 ** 2
  • p.numel()：该参数张量的元素总数
  • p.element_size()：每个元素占多少字节（FP32=4字节，INT8=1字节）
  • 两者相乘再累加 = 总字节数，/1024**2 转成 MB

量化后 element_size() 从 4 变成 1，所以内存约缩小 4 倍。

④ 保存量化模型

python 复制代码 
torch.save(quantized_model, conf.quantized_model_save_path)

⚠️ 注意 ：量化模型保存时，是保存整个模型对象（不是只存 state_dict），因为量化后的结构变了。

九、量化后的预测（predict_fun.py）

python 复制代码 
import torch
from config import Config

conf = Config()
device = conf.device
tokenizer = conf.tokenizer

# ★ 加载量化模型：整体加载，需要 weights_only=False
model = torch.load(
    "models_save/quantized_model_xxx.pt",
    map_location=device,
    weights_only=False      # 量化模型含自定义对象，必须设 False
)

def predict(data):
    text = data["text"]
    if not text.strip():
        return {"text": text, "pred_class": None}

    # 分词编码
    encoded = tokenizer.encode_plus(text, return_tensors="pt")
    input_ids = encoded["input_ids"].to(device)
    attention_mask = encoded["attention_mask"].to(device)

    # 推理
    with torch.no_grad():
        logits = model(input_ids, attention_mask)
        pred_idx = torch.argmax(logits, dim=1).item()
        pred_class = conf.class_list[pred_idx]

    return {"text": text, "pred_class": pred_class}

if __name__ == "__main__":
    result = predict({"text": "中华女子学院：本科层次仅1专业招男生"})
    print(result)   # {'text': '...', 'pred_class': 'education'}

十、运行与效果验证

10.1 运行量化

bash 复制代码 
cd 06-model-compression/bert_quantization
python bert_model_quantization.py

10.2 预期输出

复制代码 
加载数据...
加载模型...
查看量化前的模型结构=========================
BertClassifier(
  (bert): BertModel(...)
  (fc): Linear(in_features=768, out_features=10, bias=True)
)
量化后的模型=========================
BertClassifier(
  (bert): BertModel(...)
  (fc): DynamicQuantizedLinear(in_features=768, out_features=10)
)
8-bit 量化后的模型内存: 102.35 MB
保存量化模型成功！

10.3 对比效果

指标 量化前 (FP32) 量化后 (INT8) 变化
模型文件大小 409 MB ~100 MB ↓ 4 倍
推理速度 基准 更快（CPU 上提升明显） ↑
准确率 95%+ 95%+（几乎不变） ≈

⚠️ 注意：quantize_dynamic 只量化了 nn.Linear 层，BERT 的 Embedding 层仍是 FP32，所以实际压缩比不会精确到 4 倍，但效果仍然显著。

十一、常见问题（FAQ）

Q1：量化后的模型能用 load_state_dict 加载吗？

❌ 不能 。量化后结构变了（Linear → DynamicQuantizedLinear），必须整体 torch.load 加载，并且 weights_only=False。

Q2：weights_only 报错怎么办？

新版 PyTorch 默认 weights_only=True，加载含自定义对象的量化模型会失败。解决：

python 复制代码 
torch.load(path, weights_only=False)

或添加安全全局：

python 复制代码 
torch.serialization.add_safe_globals([BertClassifier])

Q3：量化后精度掉了怎么办？

  • 动态量化损失通常很小；若下降明显，可尝试静态量化 （quantize_static）
  • 或用 qconfig 精细配置量化策略

Q4：GPU 上能用量化模型吗？

动态量化主要针对 CPU 推理优化 。GPU 上建议用 torch.cuda.amp（混合精度）或 TensorRT 量化。

Q5：为什么内存缩小不到 4 倍？

因为只量化了 nn.Linear，Embedding、LayerNorm 等仍是 FP32。想进一步压缩可扩大量化层范围。

Q6：量化需要重新训练吗？

不需要。量化是"训练后压缩"（Post-Training Quantization），直接对训练好的模型操作。

十二、量化 vs 蒸馏 vs 剪枝（对比总结）

维度 量化 蒸馏 剪枝
原理 降精度（32→8位） 大模型教小模型 删冗余参数
模型结构 不变 改变（换小模型） 变瘦
代码量 极少（1行核心） 复杂（软/硬损失） 中等
压缩比 ~4 倍 几十倍 可变
精度损失 极小 略降 略降
是否需重训 不需要 需要训练学生 通常需微调

十三、总结

复制代码 
训练好的 BERT (409MB)
      ↓  torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)
量化后模型 (~100MB)
      ↓  torch.save / torch.load(weights_only=False)
部署预测（更快、更省内存、精度几乎不变）

核心一句话：

量化就是把模型参数的精度从 32 位浮点降到 8 位整数，用 PyTorch 一行代码即可完成，模型体积缩小约 4 倍，精度几乎不损失------这是工业界最常用、最省事的模型压缩手段。

附：量化速查代码

python 复制代码 
import torch

# 1. 加载训练好的模型
model = YourModel()
model.load_state_dict(torch.load("model.pt", map_location='cpu'))
model.eval()

# 2. 动态量化（核心一行）
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

# 3. 查看内存
size_mb = sum(p.numel() * p.element_size() for p in quantized_model.parameters()) / 1024**2
print(f"量化后: {size_mb:.2f} MB")

# 4. 保存（整体保存）
torch.save(quantized_model, "quantized_model.pt")

# 5. 加载预测（注意 weights_only=False）
q_model = torch.load("quantized_model.pt", map_location='cpu', weights_only=False)

祝你压缩顺利！🚀*

相关推荐
月疯9 天前
bert的架构解析
人工智能·深度学习·bert
Jialu.12 天前
模型压缩实战：BERT 量化从 390MB 到 146MB 的实践
人工智能·深度学习·bert
Kobebryant-Manba12 天前
学习Bert微调
人工智能·学习·bert
Jialu.16 天前
中文 BERT 多任务分类项目：从模型结构到训练细节
人工智能·pytorch·分类·微软·nlp·bert
Jialu.16 天前
从 102M 到 34M：BERT 到 BiLSTM 的完整知识蒸馏实现
深度学习·nlp·bert
傲笑风17 天前
【openvino】tinybert基于openvino服务化部署（四）
人工智能·python·自然语言处理·nlp·bert·openvino
Lee_jerome22 天前
python神经网络编程入门（四十四）——微调预训练 BERT 做下游任务
微调·bert·迁移学习·文本分类·预训练·小样本·冻结特征
Tbisnic1 个月前
BGE-M3 算法详解：从模型架构到三种检索方式的数学原理
算法·自然语言处理·大模型·bert·transformer·注意力机制
bulingg1 个月前
bert输入长度有限，如何处理超长文本？
人工智能·深度学习·bert
热门推荐
01国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）02GitHub 镜像站点03蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接04历年考研数学一、数学二、数学三真题试卷及答案PDF05人工智能（AI）发展史：详细里程碑06微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新072026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？08ZCode 周末送额度活动开启：3 亿 Token 免费领取09数学公理体系大全:第六章 选择公理的等价形式及证明10【一等奖版】2026 国赛 A题 药材的烘干问题