4. 量化实战：使用PyTorch进行模型后训练量化（PTQ）与量化感知训练（QAT）
聊到端侧推理，量化是绕不开的话题。我刚开始做端侧部署时，总觉得量化就是把float32转成int8，简单得很。结果第一次把量化后的模型跑在手机上，精度直接掉了5个点，被测试同学追着问了一周。嗯，从那以后我才真正开始认真研究量化这件事。
说白了，量化就是用更少的比特数来表示模型的权重和激活值。常见的做法是把FP32压缩到INT8，甚至INT4。这样做的好处很明显——模型体积变小，推理速度变快，功耗也降下来了。但代价呢？精度损失。怎么在压缩和精度之间找到平衡，这就是我们今天要聊的核心。
4.1 后训练量化（PTQ）—— 最快上手的方式
PTQ，全称Post-Training Quantization。你想想看，模型已经训练好了，我们直接拿过来做量化，不需要重新训练。这是我最常用的方法，尤其是在项目初期快速验证的时候。
PyTorch提供了两种PTQ方式：动态量化和静态量化。我分别说说。
4.1.1 动态量化
动态量化只量化权重，激活值在推理时动态计算量化参数。好处是简单，不需要校准数据。坏处是推理速度提升有限。
import torch
import torch.quantization as quant
# 加载一个预训练模型
model = torch.load('my_model.pth')
model.eval()
# 应用动态量化
quantized_model = quant.quantize_dynamic(
model,
{torch.nn.Linear, torch.nn.LSTM}, # 只量化这些层
dtype=torch.qint8
)
# 保存量化后的模型
torch.save(quantized_model.state_dict(), 'model_dynamic_quant.pth')
4.1.2 静态量化
静态量化就复杂一些了。它需要你提供一小批校准数据，让PyTorch去统计激活值的分布，然后确定量化参数。这样做的好处是推理时不需要再计算量化参数，速度更快。
import torch
import torch.quantization as quant
# 准备模型
model = MyModel()
model.eval()
# 设置量化配置
model.qconfig = quant.get_default_qconfig('fbgemm') # x86平台
# model.qconfig = quant.get_default_qconfig('qnnpack') # ARM平台
# 准备量化（插入观察者）
quant.prepare(model, inplace=True)
# 用校准数据跑一遍
with torch.no_grad():
for data in calibration_loader:
model(data)
# 转换为量化模型
quant.convert(model, inplace=True)
# 现在model已经是量化后的了
torch.save(model.state_dict(), 'model_static_quant.pth')
4.2 量化感知训练（QAT）—— 精度优先的选择
PTQ虽然快，但有时候精度损失确实让人头疼。我记得有一次做语义分割模型，PTQ后mIoU直接掉了8个点，客户直接说不行。这时候就得请出QAT了。
QAT的核心思想是在训练过程中模拟量化效果。说白了，就是在前向传播时插入伪量化节点（FakeQuantize），让模型学会适应量化带来的误差。这样训练出来的模型，量化后精度损失会小很多。
import torch
import torch.quantization as quant
# 加载预训练模型
model = MyModel()
model.load_state_dict(torch.load('pretrained.pth'))
# 切换到训练模式
model.train()
# 设置QAT配置
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
# 准备QAT
quant.prepare_qat(model, inplace=True)
# 继续训练几个epoch
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
for epoch in range(5):
for data, target in train_loader:
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
# 训练完成后，转换为量化模型
model.eval()
quant.convert(model, inplace=True)
# 保存
torch.save(model.state_dict(), 'model_qat_quant.pth')
4.3 PTQ vs QAT：怎么选？
这个问题我经常被问到。我的建议是：
| 场景 | 推荐方法 | 原因 |
|---|---|---|
| 快速验证、原型开发 | PTQ动态量化 | 不需要数据，一行代码搞定 |
| 精度要求不高、部署时间紧 | PTQ静态量化 | 速度快，精度损失可接受 |
| 精度敏感型任务（如医疗、自动驾驶） | QAT | 精度损失最小，但需要训练时间 |
| 模型已经很大，再量化怕崩 | 先PTQ试试，不行再QAT | 我习惯的做法，省时间 |
4.4 避坑指南
做量化这几年，我踩过的坑不少。挑几个典型的说说：
- 校准数据量不是越多越好。我曾经用10万张图片做校准，结果量化后的模型精度反而比用1000张的差。后来发现是因为数据分布太杂，量化参数被拉偏了。一般100-500张就够了。
- 注意量化粒度。Per-tensor量化速度快，但精度差。Per-channel量化精度好，但速度慢。我在移动端项目里一般用per-channel，因为精度更重要。
- 不要量化所有层。有些层对量化特别敏感，比如第一层和最后一层。我习惯保留这些层为FP32，其他层做INT8。PyTorch支持通过配置来指定哪些层不量化。
- QAT训练时间别太长。我见过有人QAT训练了50个epoch，结果过拟合了。一般3-5个epoch就够了，最多10个。
4.5 实战总结
量化这件事，说白了就是trade-off。你要速度还是要精度？要快速部署还是要精细调优？没有银弹，只有适合你场景的方法。
我个人习惯的流程是这样的：先用PTQ动态量化快速看看效果，如果精度损失在可接受范围内，就直接用了。如果不行，上PTQ静态量化。还不行，那就老老实实做QAT。这个流程帮我省了不少时间，也避免了一上来就做QAT那种大动干戈的情况。
最后说一句，量化不是终点，部署到端侧才是。下一节我们会聊怎么把量化后的模型部署到手机和嵌入式设备上，到时候再细说。