2. 模型量化基础

任务目标

使用PyTorch对预训练MobileNetV2进行动态范围量化（Dynamic Range Quantization），将FP32权重压缩为INT8，对比量化前后模型文件大小与ImageNet验证集上的Top-1精度变化。

环境依赖

• Python 3.8+
• PyTorch 2.0+（torch.quantization.quantize_dynamic）
• torchvision 0.15+（提供预训练权重与数据集加载）
• Pillow、numpy（图像预处理）

核心实现

以下代码完成：加载预训练MobileNetV2 → 动态量化（仅量化Linear层）→ 保存/加载量化模型 → 在100张ImageNet验证子集上对比精度与文件大小。

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import os, time, json

# ---------- 1. 加载预训练模型 ----------
model_fp32 = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
model_fp32.eval()

# ---------- 2. 动态范围量化（仅量化Linear层） ----------
model_quantized = torch.quantization.quantize_dynamic(
    model_fp32, {nn.Linear}, dtype=torch.qint8
)

# ---------- 3. 保存与文件大小对比 ----------
torch.save(model_fp32.state_dict(), "mobilenet_fp32.pth")
torch.save(model_quantized.state_dict(), "mobilenet_int8.pth")
size_fp32 = os.path.getsize("mobilenet_fp32.pth") / 1024 / 1024
size_int8 = os.path.getsize("mobilenet_int8.pth") / 1024 / 1024
print(f"FP32模型大小: {size_fp32:.2f} MB")
print(f"INT8动态量化模型大小: {size_int8:.2f} MB")
print(f"压缩比: {size_fp32 / size_int8:.2f}x")

# ---------- 4. 验证精度（ImageNet子集） ----------
# 假设已有验证集目录结构: val/<synset_id>/<image>.JPEG
# 此处使用100张图片子集作为快速验证
val_dir = "val_subset"  # 包含100张真实ImageNet验证图片
with open("imagenet_labels.json") as f:  # {文件名: 类别索引}
    label_map = json.load(f)

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def evaluate(model, device="cpu"):
    correct = 0
    total = 0
    with torch.no_grad():
        for fname, label in label_map.items():
            img = Image.open(os.path.join(val_dir, fname)).convert("RGB")
            input_tensor = transform(img).unsqueeze(0).to(device)
            output = model(input_tensor)
            pred = output.argmax(dim=1).item()
            correct += (pred == label)
            total += 1
    return correct / total

# 量化模型在CPU上推理
acc_fp32 = evaluate(model_fp32)
acc_int8 = evaluate(model_quantized)
print(f"FP32 Top-1 精度: {acc_fp32*100:.2f}%")
print(f"INT8动态量化 Top-1 精度: {acc_int8*100:.2f}%")
print(f"精度变化: {acc_int8 - acc_fp32:+.4f}")

运行与验证

1. 准备数据：从ImageNet验证集复制100张图片到val_subset/，创建imagenet_labels.json映射文件名→类别索引（0-999）。
2. 运行脚本：python quantize_demo.py
3. 通过条件：
   • 控制台输出FP32与INT8模型文件大小，压缩比约3.5~4x（因权重占主导，Linear层量化后整体压缩明显）；
   • INT8精度相对FP32下降不超过1%（动态量化通常精度损失极小）；
   • 量化模型可正常加载并完成推理，无报错。

关键说明

• 动态范围量化仅对权重做INT8映射，激活仍为FP32，因此无需校准数据集，适合快速压缩模型。
• MobileNetV2的卷积层占参数量大头，但本案例仅量化Linear层（分类头），若需进一步压缩需使用量化感知训练或静态量化。
• 若本机无GPU，上述代码默认CPU推理即可；若使用GPU，需将模型移至GPU并确保输入也在GPU。