PyTorch 动态量化（Dynamic Quantization）

PyTorch 的动态量化（Dynamic Quantization），包括原理、适用场景和具体示例。

---

1. 动态量化的概念

动态量化 (Dynamic Quantization) 是一种 后训练量化（Post-Training Quantization） 方法。

特点：

1. 只量化权重：

   • 模型的权重从 float32 → int8，减小模型存储大小。

2. 激活在推理时动态量化：

   • 前向传播时，激活会在每次计算时动态转换成 int8，再做矩阵运算。

3. 不需要重新训练：

   • 直接对训练好的模型量化即可。

4. 适合 CPU 推理：

   • 可以显著加速模型推理，同时减少内存占用。

5. 适合全连接层和 LSTM：

   • 对 Transformer、RNN、LSTM、全连接层效果好，对卷积层效果有限。

---

2. 动态量化原理

• 对于一个全连接层：

  y = x @ W.T + b

  • W：权重矩阵 → 量化成 int8
  • x：输入激活 → 保持 float32，前向计算时动态量化

• 矩阵乘法在 int8 上计算，然后转换回 float32 输出。

优势：

• 模型大小减小 3~4 倍
• 推理速度提升 2~4 倍（CPU 上明显）

---

3. PyTorch 使用示例

以你之前的 SimpleModel 为例：

import torch
import torch.nn as nn
from torch.quantization import quantize_dynamic

# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128, 96)
        self.fc2 = nn.Linear(96, 64)
        self.fc3 = nn.Linear(64, 32)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        out = self.fc3(x)
        return out

# 创建并训练好的模型（假设已经训练好）
model = SimpleModel()
# ----------------------------
# 模拟训练完成，直接量化
# ----------------------------

# 对全连接层进行动态量化
quantized_model = quantize_dynamic(
    model,                      # 原始模型
    {nn.Linear},                # 只量化 Linear 层
    dtype=torch.qint8            # 量化数据类型
)

# 保存量化后的模型
torch.save(quantized_model.state_dict(), "simple_model_dynamic_quant.pth")

# 使用量化模型进行推理
quantized_model.eval()
x_new = torch.randn(5, 128)
with torch.no_grad():
    y_pred = quantized_model(x_new)

print("动态量化模型输出形状:", y_pred.shape)  # [5, 32]

---

4. 动态量化特点总结

特性	动态量化
权重类型	int8
激活类型	float32 → 推理时动态量化
是否需要训练	不需要
适用层	Linear, LSTM, Transformer
优势	模型小，推理快，兼容 CPU
缺点	对卷积网络加速有限，精度可能略下降

---

✅ 小结：

• 动态量化：训练后直接量化 → CPU 推理加速 → 不改训练代码
• 静态量化/量化感知训练 (QAT)：需要校准或训练 → 精度更高 → GPU/CPU 均可

---