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为例：

importtorchimporttorch.nnasnnfromtorch.quantizationimportquantize_dynamic# 定义模型classSimpleModel(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()defforward(self,x):x=self.relu(self.fc1(x))x=self.relu(self.fc2(x))out=self.fc3(x)returnout# 创建并训练好的模型（假设已经训练好）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)withtorch.no_grad():y_pred=quantized_model(x_new)print("动态量化模型输出形状:",y_pred.shape)# [5, 32]

4. 动态量化特点总结

✅小结：

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