PyTorch 模型静态量化实战：ResNet-18 从 FP32 到 INT8 的 4 倍压缩与精度对比

PyTorch 模型静态量化实战ResNet-18 从 FP32 到 INT8 的 4 倍压缩与精度对比在移动端和边缘计算场景中模型大小和推理速度往往是关键瓶颈。本文将带你完整实现 ResNet-18 模型的静态量化Post-Training Quantization流程从 FP32 到 INT8 的转换过程中模型体积缩小为原来的 1/4同时保持可接受的精度损失。我们将使用 PyTorch 1.9 的量化API包含校准、量化和评估全流程。1. 量化准备与环境配置量化前需要确保 PyTorch 版本支持量化操作。推荐使用 PyTorch 1.9 及以上版本并安装必要的依赖pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113量化核心工具库导入import torch import torch.nn as nn import torch.optim as optim from torch.quantization import QuantStub, DeQuantStub, prepare, convert from torchvision.models import resnet18 from torch.utils.data import DataLoader from torchvision.datasets import CIFAR10 import torchvision.transforms as transforms数据集准备以CIFAR10为例def prepare_data(batch_size32): transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) train_set CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) val_set CIFAR10(root./data, trainFalse, downloadTrue, transformtransform) return DataLoader(train_set, batch_sizebatch_size, shuffleTrue), DataLoader(val_set, batch_sizebatch_size)2. 原始模型训练与评估首先加载预训练的 ResNet-18 并微调以适应 CIFAR1032x32 输入尺寸def train_model(model, train_loader, epochs5): criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.001, momentum0.9) model.train() for epoch in range(epochs): for inputs, labels in train_loader: optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() def evaluate(model, val_loader): correct 0 total 0 model.eval() with torch.no_grad(): for inputs, labels in val_loader: outputs model(inputs) _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() return correct / total # 加载并调整模型 model_fp32 resnet18(pretrainedTrue) model_fp32.fc nn.Linear(512, 10) # CIFAR10有10类 train_loader, val_loader prepare_data() train_model(model_fp32, train_loader) print(fFP32模型准确率: {evaluate(model_fp32, val_loader):.4f})3. 静态量化完整流程PyTorch静态量化分为三个关键步骤模型修改、校准和转换。3.1 模型量化准备首先需要修改模型以支持量化class QuantizedResNet18(nn.Module): def __init__(self, model_fp32): super().__init__() self.quant QuantStub() self.dequant DeQuantStub() self.model model_fp32 def forward(self, x): x self.quant(x) x self.model(x) x self.dequant(x) return x model_to_quantize QuantizedResNet18(model_fp32) model_to_quantize.eval()3.2 校准过程校准用于确定激活值的动态范围def calibrate_model(model, val_loader): model.eval() with torch.no_grad(): for inputs, _ in val_loader: model(inputs) # 准备量化配置 model_to_quantize.qconfig torch.quantization.get_default_qconfig(fbgemm) quantized_model prepare(model_to_quantize) # 运行校准 calibrate_model(quantized_model, val_loader)3.3 模型转换将校准后的模型转换为量化版本quantized_model convert(quantized_model) print(quantized_model)4. 量化效果评估4.1 模型大小对比def print_model_size(model, name): torch.save(model.state_dict(), temp.pth) size os.path.getsize(temp.pth) / (1024 * 1024) print(f{name}模型大小: {size:.2f}MB) os.remove(temp.pth) print_model_size(model_fp32, FP32) print_model_size(quantized_model, INT8)典型输出结果FP32模型大小: 44.7MB INT8模型大小: 11.2MB4.2 推理速度测试使用随机输入测试推理时间import time def test_inference_time(model, input_shape(1,3,32,32), iterations100): model.eval() inputs torch.rand(input_shape) start time.time() with torch.no_grad(): for _ in range(iterations): _ model(inputs) elapsed (time.time() - start) / iterations * 1000 return f{elapsed:.2f}ms print(fFP32推理时间: {test_inference_time(model_fp32)}) print(fINT8推理时间: {test_inference_time(quantized_model)})4.3 精度对比fp32_acc evaluate(model_fp32, val_loader) int8_acc evaluate(quantized_model, val_loader) print(fFP32准确率: {fp32_acc:.4f}, INT8准确率: {int8_acc:.4f}, 下降: {(fp32_acc-int8_acc):.4f})5. 量化技术细节解析5.1 对称 vs 非对称量化PyTorch默认使用对称量化其数学表示为$$ Q \text{round}(R/S) \ S \frac{\max(|R_{\min}|, |R_{\max}|)}{2^{b-1}-1} $$其中 $b$ 为量化位数INT8中b85.2 逐层量化策略PyTorch支持两种量化粒度量化类型特点适用场景逐张量整个tensor使用相同scale计算效率高逐通道每个channel单独量化精度更高5.3 量化配置选项# 自定义量化配置示例 custom_qconfig torch.quantization.QConfig( activationtorch.quantization.MinMaxObserver.with_args( qschemetorch.per_tensor_symmetric, dtypetorch.qint8 ), weighttorch.quantization.MinMaxObserver.with_args( qschemetorch.per_tensor_symmetric, dtypetorch.qint8 ) )6. 实际应用建议校准数据选择使用500-1000张具有代表性的校准图片敏感层处理第一层和最后一层对量化敏感可保持FP32精度量化感知训练当PTQ精度损失过大时考虑QAT方案# 敏感层排除量化示例 quantized_model torch.quantization.quantize_dynamic( model_fp32, {nn.Linear, nn.Conv2d}, dtypetorch.qint8, exclude_layers[fc] )7. 常见问题排查精度下降严重检查校准数据集是否具有代表性尝试调整量化位宽如INT16验证模型是否有异常大的权重值推理速度未提升确认运行环境支持INT8加速如使用FBGEMM后端检查是否真正调用了量化内核提示量化后的模型在保存和加载时需要特别注意兼容性问题建议使用PyTorch同一版本进行序列化操作