2. PyTorch动态量化实现与精度评估

任务目标

在PyTorch中实现一个面向CPU部署的动态量化流程：对训练好的浮点模型（以两层线性网络为例）应用torch.ao.quantization.quantize_dynamic，将权重从FP32量化为INT8，并对比量化前后模型在MNIST测试集上的推理精度与参数量变化，验证动态量化在仅压缩权重、不改变激活精度前提下的精度保持能力。

环境与依赖

  • Python 3.8+，PyTorch ≥ 2.0（需包含torch.ao.quantization模块）
  • torchvision（用于加载MNIST数据集）
  • 无需GPU，CPU即可完成全部训练与量化流程

完整实现代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# ---------- 1. 定义浮点模型 ----------
class TwoLayerNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# ---------- 2. 数据加载 ----------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

# ---------- 3. 训练浮点模型（2个epoch即可达到足够精度） ----------
model = TwoLayerNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(2):
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# ---------- 4. 动态量化 ----------
quantized_model = torch.ao.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

# ---------- 5. 精度评估 ----------
def evaluate(model, loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100.0 * correct / total

fp32_acc = evaluate(model, test_loader)
int8_acc = evaluate(quantized_model, test_loader)

# ---------- 6. 模型大小对比 ----------
def get_model_size_mb(model):
    torch.save(model.state_dict(), 'temp.pth')
    size = os.path.getsize('temp.pth') / 1024 / 1024
    os.remove('temp.pth')
    return size

import os
fp32_size = get_model_size_mb(model)
int8_size = get_model_size_mb(quantized_model)

print(f"FP32 模型精度: {fp32_acc:.2f}%")
print(f"INT8 动态量化模型精度: {int8_acc:.2f}%")
print(f"精度下降: {fp32_acc - int8_acc:.2f}%")
print(f"FP32 模型大小: {fp32_size:.3f} MB")
print(f"INT8 模型大小: {int8_size:.3f} MB")
print(f"压缩比: {fp32_size / int8_size:.2f}x")

运行方式

  1. 将上述代码保存为dynamic_quantization.py；
  2. 直接执行python dynamic_quantization.py；
  3. 首次运行会自动下载MNIST数据集至./data目录。

验证方法与通过标准

  • 精度指标：FP32模型在MNIST测试集上精度应≥97%（2个epoch训练可达到）；动态量化后精度下降通常不超过0.5%；
  • 压缩指标：INT8模型文件大小应约为FP32的1/4，压缩比接近4x；
  • 功能正确性：程序正常输出6行指标，且无报错；量化模型可正常完成前向推理。

关键说明

  • 动态量化仅将权重量化为INT8，激活仍以FP32计算，因此推理精度损失极小，适合LSTM、Transformer等权重主导的模型在CPU上加速；
  • 若需进一步压缩激活（静态量化），需提供校准数据并处理算子融合，不属于本案例范围；
  • 实际部署时，量化后的模型可用torch.jit.save导出为TorchScript格式，再通过LibTorch或ONNX Runtime加载。
← 上一章 📖 返回目录 下一章 →