2. PyTorch动态量化实现

任务目标

本案例将训练一个基于LSTM的语言模型（字符级），随后对其应用PyTorch动态量化（torch.quantization.quantize_dynamic），对比量化前后模型的推理速度与内存占用差异。动态量化仅将权重从FP32转为INT8，激活仍以FP32计算，特别适合LSTM这类以权重访存为主的层。

环境依赖

Python 3.8+，PyTorch 1.8+（推荐2.x），NumPy。无需额外数据集——使用随机生成的字符序列模拟训练输入，以聚焦量化流程本身。

完整实现

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

torch.manual_seed(42)

# ---------- 1. 定义LSTM语言模型 ----------
class LSTMLanguageModel(nn.Module):
    def __init__(self, vocab_size=50, embed_dim=64, hidden_dim=128, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        out, _ = self.lstm(x)
        return self.fc(out)

# ---------- 2. 初始化模型并生成模拟训练数据 ----------
model = LSTMLanguageModel()
model.train()

# 模拟一批训练数据：batch=8, seq_len=32
x = torch.randint(0, 50, (8, 32))
y = torch.randint(0, 50, (8, 32))

# 简单训练几步使权重非随机（便于量化对比）
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
for _ in range(5):
    optimizer.zero_grad()
    out = model(x)
    loss = loss_fn(out.reshape(-1, 50), y.reshape(-1))
    loss.backward()
    optimizer.step()

model.eval()

# ---------- 3. 应用动态量化 ----------
quantized_model = copy.deepcopy(model)
quantized_model = torch.quantization.quantize_dynamic(
    quantized_model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)

# ---------- 4. 推理速度与内存对比 ----------
def measure(model, input_seq, repeat=20):
    # 预热
    with torch.no_grad():
        model(input_seq)
    # 计时
    torch.cuda.synchronize() if torch.cuda.is_available() else None
    start = time.perf_counter()
    with torch.no_grad():
        for _ in range(repeat):
            model(input_seq)
    elapsed = (time.perf_counter() - start) / repeat
    # 内存（模型参数+缓冲区）
    mem_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    mem_bytes += sum(b.numel() * b.element_size() for b in model.buffers())
    return elapsed * 1000, mem_bytes / 1024  # ms, KB

test_input = torch.randint(0, 50, (1, 32))
fp32_time, fp32_mem = measure(model, test_input)
int8_time, int8_mem = measure(quantized_model, test_input)

# ---------- 5. 验证输出一致性 ----------
with torch.no_grad():
    out_fp32 = model(test_input)
    out_int8 = quantized_model(test_input)
max_diff = (out_fp32 - out_int8).abs().max().item()

print(f"FP32 推理耗时: {fp32_time:.3f} ms/样本")
print(f"INT8 推理耗时: {int8_time:.3f} ms/样本")
print(f"FP32 模型内存: {fp32_mem:.2f} KB")
print(f"INT8 模型内存: {int8_mem:.2f} KB")
print(f"输出最大差异: {max_diff:.6f}")

运行方式

将上述代码保存为 dynamic_quant_lstm.py，在终端执行：

python dynamic_quant_lstm.py

验证方法与判定标准

• 功能正确性：程序无报错正常退出，且打印的输出最大差异小于0.1（INT8量化对输出影响极小，说明量化模型功能保持）。
• 内存减少：INT8 模型内存应明显小于FP32 模型内存（通常降低约40%~50%，因LSTM与Linear权重占主导）。
• 速度提升：INT8 推理耗时应低于FP32 推理耗时。若差异不明显，可增大repeat次数或使用更大hidden_dim（如256）以放大差异。

预期输出示例

FP32 推理耗时: 2.847 ms/样本
INT8 推理耗时: 1.923 ms/样本
FP32 模型内存: 214.34 KB
INT8 模型内存: 112.18 KB
输出最大差异: 0.003821

注：以上数值仅为逻辑一致的示例，实际值因CPU型号、PyTorch版本而异，但内存下降比例与输出差异量级应保持一致。

常见问题

• 速度不升反降：输入序列过短或模型过小导致量化开销占比高。增大hidden_dim至256或seq_len至64后重试。
• 量化后精度骤降：确认模型已调用eval()；若仍异常，检查是否误量化了Embedding层（动态量化仅支持LSTM和Linear）。
• 内存未明显减少：确认模型权重确实以INT8存储——打印quantized_model.lstm.weight_ih_l0.dtype应显示torch.qint8。