4. 量化基础
任务目标
对训练好的PyTorch模型应用动态量化，将权重从FP32转为INT8，并量化后模型在CPU上的推理速度与精度变化。本案例使用公开的ResNet18模型在CIFAR-10子集上完成闭环验证。
环境依赖
Python 3.8+，PyTorch ≥ 1.13（需包含torch.ao.quantization模块），torchvision，numpy。CPU即可完成全部实验。
完整实现
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import time
import copy
# ---------- 1. 准备数据与模型 ----------
transform = 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)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)
model = torchvision.models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 10) # CIFAR-10分类
# 微调最后一个全连接层（仅演示用，保持其他层预训练权重）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
model.train()
for epoch in range(1): # 1个epoch足够演示
for inputs, labels in trainloader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
model.eval()
# ---------- 2. 评估函数 ----------
def evaluate(model, dataloader):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
return 100.0 * correct / total
def measure_inference_time(model, dataloader, num_batches=20):
model.eval()
model.to('cpu')
times = []
with torch.no_grad():
for i, (inputs, _) in enumerate(dataloader):
if i >= num_batches:
break
inputs = inputs.to('cpu')
start = time.time()
_ = model(inputs)
times.append(time.time() - start)
return sum(times) / len(times) * 1000 # 返回毫秒
# ---------- 3. 动态量化 ----------
model_cpu = copy.deepcopy(model).to('cpu')
quantized_model = torch.ao.quantization.quantize_dynamic(
model_cpu, {nn.Linear}, dtype=torch.qint8
)
# ---------- 4. 精度与速度对比 ----------
# 精度（在CPU上评估）
model_cpu.eval()
fp32_acc = evaluate(model_cpu, testloader)
int8_acc = evaluate(quantized_model, testloader)
# 推理速度（CPU，毫秒/批）
fp32_time = measure_inference_time(model_cpu, testloader)
int8_time = measure_inference_time(quantized_model, testloader)
# 模型大小对比
def get_model_size_mb(model):
param_size = sum(p.numel() * p.element_size() for p in model.parameters())
buffer_size = sum(b.numel() * b.element_size() for b in model.buffers())
return (param_size + buffer_size) / (1024 * 1024)
fp32_size = get_model_size_mb(model_cpu)
int8_size = get_model_size_mb(quantized_model)
print(f"FP32 精度: {fp32_acc:.2f}% | INT8 精度: {int8_acc:.2f}% | 精度差: {fp32_acc - int8_acc:.2f}%")
print(f"FP32 推理: {fp32_time:.2f} ms/批 | INT8 推理: {int8_time:.2f} ms/批 | 加速比: {fp32_time/int8_time:.2f}x")
print(f"FP32 大小: {fp32_size:.2f} MB | INT8 大小: {int8_size:.2f} MB | 压缩比: {fp32_size/int8_size:.2f}x")运行方式
- 将上述代码保存为
dynamic_quantization.py； - 首次运行会自动下载CIFAR-10数据集（约170MB）和预训练ResNet18权重；
- 执行
python dynamic_quantization.py，等待训练与量化完成。
验证方法与判定标准
程序输出三组对比数据，满足以下条件即判定量化成功：
- 精度损失：INT8精度相对FP32下降不超过2个百分点（本案例仅微调最后一层，预训练特征提取能力保留，通常损失在0.5~1.5%之间）；
- 推理加速：INT8推理时间明显低于FP32（CPU上通常有1.5~3倍加速）；
- 模型压缩：INT8模型大小约为FP32的25%~30%（因仅量化Linear层，卷积层仍为FP32）。
若精度下降过大，可增加微调epoch数；若加速不明显，确认已设置model.eval()且使用torch.no_grad()，并确保线程数未被限制。
常见问题
- 量化后精度大幅下降：检查是否对全连接层进行了充分微调；动态量化对权重分布敏感的层影响较大，可尝试仅量化部分层。
- 无加速效果：确认在CPU上运行且输入批大小足够大（≥32），小批量时线程调度开销可能抵消量化收益。
- 模块导入错误：PyTorch 1.13以下版本使用
torch.quantization替代torch.ao.quantization。