4. 动态量化实现

任务目标

使用 PyTorch 对端侧语音模型（以 LSTM 声学特征提取网络为例）实施动态量化，对比量化前后模型文件大小与 CPU 推理速度，验证动态量化在保持可用精度的前提下对模型体积和延迟的优化效果。

环境依赖

• Python 3.8+
• PyTorch ≥ 1.8（torch.quantization 模块）
• torchaudio（用于生成测试输入，仅需基础波形读取功能）

实现代码

import torch
import torch.nn as nn
import time
import copy

# ---------- 1. 定义语音特征提取模型（LSTM + 线性层） ----------
class SpeechFeatureNet(nn.Module):
    def __init__(self, input_size=40, hidden_size=128, num_layers=2, num_classes=10):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_size, 64),
            nn.ReLU(),
            nn.Linear(64, num_classes)
        )

    def forward(self, x):
        # x: (batch, time, input_size)
        out, _ = self.lstm(x)
        # 取最后一个时间步输出
        out = out[:, -1, :]
        return self.classifier(out)

# ---------- 2. 初始化模型并生成确定性输入 ----------
torch.manual_seed(42)
model = SpeechFeatureNet()
model.eval()

# 模拟 1 秒语音，16kHz 采样率，40 维 MFCC 特征，帧移 10ms → 100 帧
batch_size = 1
seq_len = 100
dummy_input = torch.randn(batch_size, seq_len, 40)

# ---------- 3. 动态量化 ----------
quantized_model = torch.quantization.quantize_dynamic(
    model,  # 原始模型
    {nn.LSTM, nn.Linear},  # 量化层类型
    dtype=torch.qint8  # 权重量化为 8 位整数
)

# ---------- 4. 模型大小对比 ----------
def get_model_size_mb(model):
    torch.save(model.state_dict(), "tmp_model.pth")
    size = os.path.getsize("tmp_model.pth") / (1024 * 1024)
    return size

import os
orig_size = get_model_size_mb(model)
quant_size = get_model_size_mb(quantized_model)
print(f"原始模型大小: {orig_size:.4f} MB")
print(f"动态量化模型大小: {quant_size:.4f} MB")
print(f"体积缩减比: {orig_size / quant_size:.2f}x")

# ---------- 5. CPU 推理速度对比 ----------
def benchmark_inference(model, input_tensor, num_runs=50):
    # 预热
    with torch.no_grad():
        for _ in range(5):
            model(input_tensor)
    # 计时
    torch.cuda.synchronize() if torch.cuda.is_available() else None
    start = time.time()
    with torch.no_grad():
        for _ in range(num_runs):
            model(input_tensor)
    elapsed = time.time() - start
    return elapsed / num_runs * 1000  # 返回单次毫秒

orig_latency = benchmark_inference(model, dummy_input)
quant_latency = benchmark_inference(quantized_model, dummy_input)
print(f"原始模型平均推理延迟: {orig_latency:.2f} ms")
print(f"动态量化模型平均推理延迟: {quant_latency:.2f} ms")
print(f"加速比: {orig_latency / quant_latency:.2f}x")

# ---------- 6. 精度一致性检查（输出差异） ----------
with torch.no_grad():
    orig_output = model(dummy_input)
    quant_output = quantized_model(dummy_input)

# 计算最大误差
max_diff = (orig_output - quant_output).abs().max().item()
print(f"量化前后输出最大误差: {max_diff:.6f}")

# 验证分类结果是否一致
orig_pred = orig_output.argmax(dim=1)
quant_pred = quant_output.argmax(dim=1)
print(f"原始模型预测类别: {orig_pred.item()}, 量化模型预测类别: {quant_pred.item()}")
print(f"预测一致性: {'通过' if orig_pred.item() == quant_pred.item() else '失败'}")

运行方式

# 安装依赖
pip install torch torchaudio

# 运行脚本
python dynamic_quantization_demo.py

验证方法与判定标准

• 模型大小：量化后模型文件应显著小于原始模型，通常缩减至原来的 1/3～1/4（体积缩减比 > 2.5x 即视为有效）。
• 推理速度：量化后模型在 CPU 上的单次推理延迟应低于原始模型，加速比 > 1.2x 即表明动态量化带来性能提升。
• 精度保持：量化前后模型对同一输入的预测类别应保持一致（预测一致性为“通过”），最大误差通常小于 0.1。

以上三项指标全部满足，即判定动态量化实现成功。若加速比不足，可增加 num_runs 至 200 次以降低计时噪声；若精度差异过大，可检查模型是否处于 eval() 模式（动态量化要求模型在推理模式下进行）。