📚 课程目录第3章量化感知训练入门

3. 量化感知训练入门

任务目标

在PyTorch中实现量化感知训练（QAT），对预训练模型插入伪量化节点，通过微调使模型权重和激活分布适应INT8量化，从而恢复因量化带来的精度损失。本案例以ResNet-18在CIFAR-10上的分类任务为例，完成从浮点模型到QAT微调的完整闭环。

环境依赖

需要PyTorch ≥ 1.8（推荐1.13+），torchvision，以及CIFAR-10数据集（首次运行自动下载）。无需额外安装量化库，PyTorch原生torch.ao.quantization模块即提供完整QAT支持。

关键完整代码


import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.ao.quantization import QConfig, default_qat_qconfig
from torch.ao.quantization.quantize_fx import prepare_qat_fx, convert_fx

# ---------- 1. 数据加载 ----------
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=256, shuffle=False, num_workers=2)

# ---------- 2. 浮点模型准备 ----------
model = torchvision.models.resnet18(num_classes=10)
model.load_state_dict(torch.load('resnet18_cifar10_fp32.pth'))  # 预训练浮点权重
model.train()

# ---------- 3. 配置QAT ----------
model.qconfig = default_qat_qconfig  # 默认8位对称/非对称量化配置
model = prepare_qat_fx(model, {'': torch.ao.quantization.default_qconfig})

# ---------- 4. QAT微调 ----------
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4)
criterion = nn.CrossEntropyLoss()
epochs = 5

for epoch in range(epochs):
    running_loss = 0.0
    for images, labels in trainloader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(trainloader):.4f}')

# ---------- 5. 转换并评估 ----------
model.eval()
model = convert_fx(model)  # 将伪量化节点转为真实INT8量化

correct = 0
total = 0
with torch.no_grad():
    for images, labels in testloader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print(f'QAT模型测试准确率: {100 * correct / total:.2f}%')

# ---------- 6. 保存量化模型 ----------
torch.save(model.state_dict(), 'resnet18_cifar10_qat.pth')

运行与验证

前置条件：需要先准备浮点预训练权重resnet18_cifar10_fp32.pth。若没有该权重，可先运行以下代码训练一个浮点模型（约10轮即可达到85%+准确率）：


# 浮点模型快速训练（仅需运行一次）
model = torchvision.models.resnet18(num_classes=10)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
for epoch in range(10):
    model.train()
    for images, labels in trainloader:
        optimizer.zero_grad()
        loss = criterion(model(images), labels)
        loss.backward()
        optimizer.step()
    print(f'FP32 Epoch {epoch+1} done')
torch.save(model.state_dict(), 'resnet18_cifar10_fp32.pth')

验证方法：运行QAT代码后，观察两点：

训练损失应逐轮下降，说明模型正在适应量化噪声；
最终打印的QAT模型测试准确率应与浮点模型准确率差距在2%以内（浮点模型通常约88-90%，QAT后应不低于86%）。

若QAT后准确率显著低于浮点模型（差距超过5%），可尝试增加微调轮数至10轮或降低学习率至0.0005。

关键说明

本案例使用FX图模式量化（prepare_qat_fx），它自动在卷积和全连接层后插入伪量化节点，模拟INT8量化的舍入误差。微调过程中，这些伪量化节点的scale和zero_point会随数据分布动态更新，使模型权重逐渐适应低精度表示。convert_fx将训练好的模型转换为真正的INT8推理模型，此时权重以INT8存储，激活在推理时动态量化。若需进一步加速，可配合torch.jit或onnxruntime部署。