2. 量化感知训练模拟
任务说明
本案例目标是在训练过程中模拟INT8量化误差，通过量化感知训练（QAT）微调模型，使量化后的INT8模型精度接近原始FP32模型。案例使用PyTorch 2.0+的torch.ao.quantization接口，在MNIST分类任务上完成从FP32训练到QAT微调、再到INT8静态量化的完整闭环。
环境依赖
- Python 3.9+，PyTorch 2.0+（需包含
torch.ao.quantization模块） - torchvision，MNIST数据集（首次运行自动下载）
完整实现
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.ao.quantization import QConfig, MinMaxObserver, MovingAverageMinMaxObserver
from torch.ao.quantization.quantize_fx import prepare_qat_fx, convert_fx
# ---------- 1. 定义模型（FP32训练用） ----------
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
self.bn1 = nn.BatchNorm2d(8)
self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
self.bn2 = nn.BatchNorm2d(16)
self.fc = nn.Linear(16 * 7 * 7, 10)
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.max_pool2d(x, 2)
x = F.relu(self.bn2(self.conv2(x)))
x = F.max_pool2d(x, 2)
x = x.view(x.size(0), -1)
return self.fc(x)
# ---------- 2. 训练/评估工具 ----------
def train_one_epoch(model, loader, opt, device):
model.train()
for data, target in loader:
data, target = data.to(device), target.to(device)
opt.zero_grad()
out = model(data)
loss = F.cross_entropy(out, target)
loss.backward()
opt.step()
def evaluate(model, loader, device):
model.eval()
correct = 0
with torch.no_grad():
for data, target in loader:
data, target = data.to(device), target.to(device)
out = model(data)
pred = out.argmax(dim=1)
correct += (pred == target).sum().item()
return correct / len(loader.dataset)
# ---------- 3. 数据加载 ----------
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_ds = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_ds = datasets.MNIST('./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=256)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ---------- 4. 阶段一：FP32模型训练 ----------
fp32_model = Net().to(device)
opt = torch.optim.Adam(fp32_model.parameters(), lr=1e-3)
for epoch in range(3):
train_one_epoch(fp32_model, train_loader, opt, device)
fp32_acc = evaluate(fp32_model, test_loader, device)
print(f'FP32 accuracy: {fp32_acc:.4f}')
# ---------- 5. 阶段二：QAT准备（模拟量化误差） ----------
qat_model = Net().to('cpu')
qat_model.load_state_dict(fp32_model.state_dict())
qat_model.train()
# 配置fake quant：权重用per-tensor对称，激活用per-tensor非对称
qconfig = QConfig(
activation=MovingAverageMinMaxObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine),
weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric)
)
qat_model = prepare_qat_fx(qat_model, {'': qconfig})
# QAT微调（2个epoch，低学习率）
opt_qat = torch.optim.Adam(qat_model.parameters(), lr=1e-4)
for epoch in range(2):
train_one_epoch(qat_model, train_loader, opt_qat, 'cpu')
# ---------- 6. 阶段三：转换为INT8静态量化模型 ----------
qat_model.eval()
int8_model = convert_fx(qat_model)
# 静态量化需要校准：用训练集子集统计激活范围
calib_loader = DataLoader(train_ds, batch_size=256, shuffle=True)
with torch.no_grad():
for i, (data, _) in enumerate(calib_loader):
int8_model(data)
if i >= 4: # 约1000张图校准
break
# ---------- 7. 验证：INT8精度对比 ----------
int8_acc = evaluate(int8_model, test_loader, 'cpu')
print(f'INT8 (QAT) accuracy: {int8_acc:.4f}')
print(f'Accuracy gap: {fp32_acc - int8_acc:.4f}')运行方式
python qat_mnist.py验证方法与通过标准
- 程序依次输出FP32精度、QAT后INT8精度及两者差值。
- 通过标准：
Accuracy gap ≤ 0.01（即INT8精度与FP32差距在1%以内），且INT8精度不低于0.98。 - 若直接对FP32模型做静态量化（不做QAT），精度通常会降至0.95以下；本案例通过QAT微调使INT8精度回升至接近FP32，验证了量化感知训练的有效性。
关键说明
prepare_qat_fx在模型中插入fake quant节点，前向传播时模拟INT8量化的舍入误差，但反向传播仍使用浮点梯度，使模型权重适应量化噪声。- QAT微调使用较低学习率（1e-4），避免破坏已收敛的FP32权重。
- 静态量化前的校准步骤统计激活值范围，是INT8推理精度的重要保障。
- 若需部署到实际端侧，可进一步用
torch.jit.script导出为TorchScript格式，或通过ONNX导出到端侧推理引擎。