PyTorch静态量化：从原理到实战
大家好，我是你们的老朋友。今天我们来聊聊PyTorch静态量化。说实话，这玩意儿我刚开始接触的时候也头大，什么FakeQuantize、校准、转换...一堆术语砸过来。但后来在项目中真正用起来，才发现它其实没那么玄乎。
静态量化，说白了就是把模型里的浮点数运算，变成整数运算。为什么要这么做？因为整数运算快啊，内存占用也小。我有个项目，模型在GPU上跑得挺欢，但一部署到手机端就卡成PPT。后来做了静态量化，体积缩小了4倍，速度提升了2倍多。嗯，真香。
核心思想：训练时模拟量化效果，推理时真正执行量化计算。
FakeQuantize原理：先骗骗模型
FakeQuantize，直译就是“假量化”。为什么叫“假”？因为它在训练时假装做了量化，但实际上还是用浮点数在算。
你想想看，如果直接把模型从浮点变成整点，精度肯定会掉。FakeQuantize的作用就是让模型提前适应这种精度损失。它在前向传播时，把浮点数值映射到离散的整点值上，再映射回浮点。这样模型就知道：“哦，原来我以后要这样干活。”
具体怎么映射？公式很简单：
scale = (max_val - min_val) / (quant_max - quant_min)
zero_point = quant_min - min_val / scale
quantized_val = round(float_val / scale) + zero_point
我在项目中遇到过一个问题：scale和zero_point怎么选？ 选不好，精度直接崩。后来我总结了一个经验：先跑一小批数据，统计出激活值的分布范围，再确定scale。别偷懒用默认值，血的教训。
小技巧：FakeQuantize的scale和zero_point是可学习的参数。如果你用QAT（量化感知训练），它们会在训练过程中自动调整。但静态量化里，它们是固定的。
torch.quantization模块介绍
PyTorch把量化相关的工具都放在了torch.quantization模块里。我刚开始用的时候，觉得这模块的API设计有点绕。但用多了就发现，它其实就干三件事：
- 准备模型：把普通模型变成可量化的形式
- 插入量化节点：在关键位置放上FakeQuantize
- 转换模型：把假量化变成真量化
常用的类有这些：
| 类名 | 作用 |
|---|---|
QuantStub |
标记输入需要量化 |
DeQuantStub |
标记输出需要反量化 |
torch.quantization.prepare |
插入FakeQuantize节点 |
torch.quantization.convert |
将FakeQuantize替换为真正的量化算子 |
我个人习惯是：先写一个简单的测试模型，跑通整个流程，再套用到复杂模型上。这样能快速定位问题。
准备模型：给模型“打疫苗”
准备模型这一步，说白了就是告诉PyTorch：“这个模型我要量化了，你做好准备”。
怎么做？主要有两种方式：
- 手动修改模型定义：在模型的前后加上
QuantStub和DeQuantStub - 使用
torch.quantization.QuantWrapper：自动包裹模型
我推荐第二种，省事。但要注意：如果你的模型里有自定义层，QuantWrapper可能不识别。这时候就得手动来了。
import torch
import torch.nn as nn
from torch.quantization import QuantStub, DeQuantStub
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.quant = QuantStub()
self.conv = nn.Conv2d(3, 64, 3)
self.relu = nn.ReLU()
self.dequant = DeQuantStub()
def forward(self, x):
x = self.quant(x)
x = self.conv(x)
x = self.relu(x)
x = self.dequant(x)
return x
注意：QuantStub和DeQuantStub必须成对出现。我曾经漏掉一个，结果模型推理结果全错，排查了半天才发现。
插入量化节点：FakeQuantize就位
这一步，PyTorch会自动帮我们做。调用torch.quantization.prepare，它会遍历模型，在每个需要量化的层前后插入FakeQuantize节点。
哪些层需要量化？一般来说：
- 卷积层：权重和激活值都要量化
- 全连接层：同上
- 激活函数：只量化输出
- BN层：通常不量化，或者折叠到卷积里
我记得有一次，模型量化后精度掉得厉害。后来发现是BN层没有折叠。BN层在推理时其实可以合并到卷积里，但PyTorch默认不帮你做。你得手动调用torch.quantization.fuse_modules。
model = MyModel()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model = torch.quantization.fuse_modules(model, [['conv', 'relu']])
model = torch.quantization.prepare(model)
避坑指南：我曾经在fuse_modules时写错了层名，结果模型直接报错。建议先打印模型结构，确认层名后再写。
校准（Calibration）：让模型“见见世面”
校准，就是用一小批真实数据跑一遍模型，统计出每层激活值的分布。这样FakeQuantize就能知道scale和zero_point该设多少。
校准数据不需要太多，几百张图片就够了。但数据分布要尽量接近真实场景。我有个项目，校准数据全是白天拍的，结果模型在夜间场景下精度暴跌。后来换了包含夜间数据的校准集，问题就解决了。
model.eval()
with torch.no_grad():
for images, _ in calibration_loader:
model(images)
校准完成后，FakeQuantize里的scale和zero_point就固定下来了。这时候可以调用torch.quantization.convert，把模型真正转换成量化版本。
关键点：校准数据不能包含标签，只需要输入数据。模型只做前向推理，不计算梯度。
转换量化模型：假戏真做
最后一步，convert。它会把FakeQuantize节点替换成真正的量化算子，比如torch.quantized.Conv2d。这时候模型就变成了一个纯整数运算的模型。
model = torch.quantization.convert(model)
转换后的模型，权重和激活值都是torch.qint8类型。你可以用model.state_dict()查看，会发现权重的数据类型变了。
我建议转换后立即做一次精度验证。用测试集跑一遍，对比量化前后的精度差异。一般来说，精度损失在1%以内都是可以接受的。如果超过2%，就得考虑调整量化配置了。
注意：转换后的模型不能再训练。因为量化算子的梯度是近似计算的，训练会不稳定。如果你需要训练，请使用QAT（量化感知训练）。
知识体系总览
为了让你更直观地理解整个流程，我画了一张图：
这张图把整个流程串起来了。你从左边开始，一步步走到右边，就完成了静态量化。每个步骤都有对应的PyTorch API，代码量其实很少。
好了，这一章的内容就到这里。静态量化是模型部署的利器，但也不是万能的。如果你的模型对精度极其敏感，或者有特殊算子，可能需要考虑其他量化方案。不过对于大多数CNN模型，静态量化已经足够用了。
记住：先跑通流程，再优化细节。别一上来就想着调参，先把整个链路跑通，后面的事就好办了。