3. 量化感知训练入门

任务目标

在CIFAR-10数据集上，使用PyTorch搭建一个简易CNN，通过torch.ao.quantization的FakeQuantize模块模拟量化误差，完成量化感知训练（QAT）。最终对比普通训练与QAT训练后的模型在验证集上的Top-1准确率，观察伪量化操作对精度的影响。

环境与依赖

Python 3.8+，PyTorch 1.13+（需包含torch.ao.quantization）
torchvision（用于CIFAR-10数据集加载）
无需GPU，CPU即可完成训练（为缩短时间，仅训练5个epoch）

关键实现说明

核心思路是在普通CNN的卷积层后插入FakeQuantize模块，使前向传播时权重和激活被量化到8bit再反量化回浮点，从而模拟推理时的量化误差。训练时保持浮点梯度更新，验证时关闭伪量化以对比真实精度。代码中通过QATNet与PlainNet两个结构完全相同的网络实现对照实验。

完整代码

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.ao.quantization import FakeQuantize, default_fake_quant
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# ---------- 基础CNN模块 ----------
class ConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch, use_fq=False):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.bn = nn.BatchNorm2d(out_ch)
        self.use_fq = use_fq
        if use_fq:
            self.fq_weight = FakeQuantize(default_fake_quant, quant_min=-128, quant_max=127)
            self.fq_act = FakeQuantize(default_fake_quant, quant_min=0, quant_max=127)

    def forward(self, x):
        if self.use_fq:
            w = self.fq_weight(self.conv.weight)
            x = F.conv2d(x, w, self.conv.bias, padding=1)
        else:
            x = self.conv(x)
        x = self.bn(x)
        x = F.relu(x)
        if self.use_fq:
            x = self.fq_act(x)
        return x

# ---------- 两种网络：普通 vs QAT ----------
class PlainNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.block1 = ConvBlock(3, 32, use_fq=False)
        self.block2 = ConvBlock(32, 64, use_fq=False)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(64 * 8 * 8, 10)

    def forward(self, x):
        x = self.pool(self.block1(x))
        x = self.pool(self.block2(x))
        x = x.view(x.size(0), -1)
        return self.fc(x)

class QATNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.block1 = ConvBlock(3, 32, use_fq=True)
        self.block2 = ConvBlock(32, 64, use_fq=True)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(64 * 8 * 8, 10)

    def forward(self, x):
        x = self.pool(self.block1(x))
        x = self.pool(self.block2(x))
        x = x.view(x.size(0), -1)
        return self.fc(x)

# ---------- 训练与验证 ----------
def train_one_epoch(model, loader, optimizer, device):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for data, target in loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        out = model(data)
        loss = F.cross_entropy(out, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        pred = out.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
    return total_loss / len(loader), 100.0 * correct / total

@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    correct, total = 0, 0
    for data, target in loader:
        data, target = data.to(device), target.to(device)
        out = model(data)
        pred = out.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
    return 100.0 * correct / total

# ---------- 主流程 ----------
def main():
    device = torch.device("cpu")
    batch_size = 128
    epochs = 5

    # 数据加载（为演示速度，仅用5000个训练样本）
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])
    full_train = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
    train_subset = torch.utils.data.Subset(full_train, range(5000))
    train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(datasets.CIFAR10(root="./data", train=False, download=True, transform=transform),
                            batch_size=batch_size, shuffle=False)

    # 训练普通网络
    plain_net = PlainNet().to(device)
    opt = optim.SGD(plain_net.parameters(), lr=0.01, momentum=0.9)
    for ep in range(epochs):
        train_loss, train_acc = train_one_epoch(plain_net, train_loader, opt, device)
        val_acc = evaluate(plain_net, val_loader, device)
        print(f"Plain Epoch {ep+1}: loss={train_loss:.4f} train_acc={train_acc:.2f}% val_acc={val_acc:.2f}%")

    # 训练QAT网络
    qat_net = QATNet().to(device)
    opt = optim.SGD(qat_net.parameters(), lr=0.01, momentum=0.9)
    for ep in range(epochs):
        train_loss, train_acc = train_one_epoch(qat_net, train_loader, opt, device)
        val_acc = evaluate(qat_net, val_loader, device)
        print(f"QAT Epoch {ep+1}: loss={train_loss:.4f} train_acc={train_acc:.2f}% val_acc={val_acc:.2f}%")

    # 最终对比
    plain_acc = evaluate(plain_net, val_loader, device)
    qat_acc = evaluate(qat_net, val_loader, device)
    print(f"\n最终验证集准确率对比:")
    print(f"普通网络: {plain_acc:.2f}%")
    print(f"QAT网络:  {qat_acc:.2f}%")
    print(f"精度差异: {qat_acc - plain_acc:+.2f}%")

if __name__ == "__main__":
    main()

运行方式

python qat_cifar10.py

首次运行会自动下载CIFAR-10数据集（约170MB）到./data目录。训练在CPU上约需5～10分钟（取决于机器性能）。

验证方法与判定标准

观察输出：程序会打印每个epoch的训练损失、训练准确率以及验证准确率。
判定通过条件：程序正常结束且输出最终对比行。QAT网络与普通网络的验证准确率差异通常在±3%以内（由于训练样本较少，波动可能稍大）。若QAT精度不低于普通网络5个百分点以上，说明伪量化模块工作正常，训练过程稳定收敛。
附加检查：若将QATNet中use_fq改为False，其精度应与PlainNet基本一致，可验证网络结构本身无差异。

常见问题

下载失败：可手动下载CIFAR-10到./data/cifar-10-batches-py，或更换网络环境重试。
精度差异过大：检查是否误将FakeQuantize的quant_min/max设置错误，或学习率过高导致不收敛。可尝试降低学习率至0.005。
伪量化不生效：确认PyTorch版本≥1.13，且FakeQuantize在torch.ao.quantization路径下可正常导入。

📚 课程目录