3. 量化感知训练入门
任务目标
对一个小型语音识别模型（基于音素分类的MLP）实施量化感知训练（QAT），对比量化前后及QAT微调后的精度变化，并理解对称量化中scale与zero_point的计算方式。
环境与依赖
torch==2.1.0
torchaudio==2.1.0
数据集与模型
使用torchaudio.datasets.SPEECHCOMMANDS的“yes/no”二分类子集，将音频预处理为13维MFCC特征。模型为两层MLP（13→64→2），参数量约1K，适合在CPU上快速完成QAT演示。
关键完整代码
import torch
import torch.nn as nn
import torchaudio
import torchaudio.transforms as T
from torch.utils.data import Dataset, DataLoader
# ---------- 1. 数据准备 ----------
class YesNoDataset(Dataset):
def __init__(self, split='training'):
self.data = torchaudio.datasets.SPEECHCOMMANDS(root='./data', download=True)
self.mfcc = T.MFCC(sample_rate=16000, n_mfcc=13)
self.labels = {'yes': 0, 'no': 1}
self.samples = []
for wav, sr, label, *_ in self.data:
if label in self.labels and wav.shape[1] > 8000:
feat = self.mfcc(wav).mean(dim=2).squeeze(0) # [13]
self.samples.append((feat, self.labels[label]))
# 简单按顺序切分，实际使用应随机打乱
n = len(self.samples)
split_idx = int(n * 0.8)
if split == 'training':
self.samples = self.samples[:split_idx]
else:
self.samples = self.samples[split_idx:]
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
return self.samples[idx]
# ---------- 2. 模型定义 ----------
class SmallMLP(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(13, 64)
self.fc2 = nn.Linear(64, 2)
self.relu = nn.ReLU()
def forward(self, x):
x = self.relu(self.fc1(x))
return self.fc2(x)
# ---------- 3. 训练/评估函数 ----------
def train_one_epoch(model, loader, opt, criterion):
model.train()
total_loss = 0
for x, y in loader:
opt.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.backward()
opt.step()
total_loss += loss.item()
return total_loss / len(loader)
def evaluate(model, loader):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for x, y in loader:
out = model(x)
pred = out.argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
# ---------- 4. 量化感知训练主流程 ----------
def main():
torch.manual_seed(0)
train_ds = YesNoDataset('training')
test_ds = YesNoDataset('testing')
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=32)
# 4.1 训练浮点基线模型
model_fp = SmallMLP()
opt = torch.optim.Adam(model_fp.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
for epoch in range(5):
train_one_epoch(model_fp, train_loader, opt, criterion)
acc_fp = evaluate(model_fp, test_loader)
print(f'浮点模型精度: {acc_fp:.4f}')
# 4.2 直接后训练量化（PTQ）对比
model_ptq = SmallMLP()
model_ptq.load_state_dict(model_fp.state_dict())
model_ptq.eval()
model_ptq.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model_ptq, inplace=True)
# 用少量训练数据校准
with torch.no_grad():
for x, _ in list(train_loader)[:5]:
model_ptq(x)
torch.quantization.convert(model_ptq, inplace=True)
acc_ptq = evaluate(model_ptq, test_loader)
print(f'PTQ量化后精度: {acc_ptq:.4f}')
# 4.3 量化感知训练（QAT）
model_qat = SmallMLP()
model_qat.load_state_dict(model_fp.state_dict())
model_qat.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(model_qat, inplace=True)
opt_qat = torch.optim.Adam(model_qat.parameters(), lr=1e-4)
for epoch in range(5):
train_one_epoch(model_qat, train_loader, opt_qat, criterion)
model_qat.eval()
torch.quantization.convert(model_qat, inplace=True)
acc_qat = evaluate(model_qat, test_loader)
print(f'QAT量化后精度: {acc_qat:.4f}')
# 4.4 观察量化参数（取fc1层）
fc1 = model_qat.fc1
print(f'fc1 weight scale: {fc1.weight().q_scale():.6f}')
print(f'fc1 weight zero_point: {fc1.weight().q_zero_point()}')
if __name__ == '__main__':
main()
运行方式
python qat_demo.py
验证方法与判定标准
程序会依次输出三个精度值。判定通过的参考标准：
- 浮点模型精度 > 0.90
- PTQ精度较浮点下降不超过5个百分点
- QAT精度不低于PTQ精度，且通常能恢复部分精度损失
- 打印的scale为正值（约0.01～0.05），zero_point为0（对称量化）
若QAT精度显著低于PTQ，可增加QAT训练轮次至10轮或调低学习率至1e-5。
量化参数解读
对称量化公式为 q = round(x / scale)，其中 scale = max_abs(x) / 127。代码中打印的q_scale()即该层的权重scale，q_zero_point()为0表示使用对称量化（无零点偏移）。QAT通过在训练中模拟量化误差（伪量化），使模型权重适应量化噪声，从而减少精度损失。