PyTorch 的动态量化(Dynamic Quantization),包括原理、适用场景和具体示例。


1. 动态量化的概念

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

特点:

  1. 只量化权重

    • 模型的权重从 float32int8,减小模型存储大小。
  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 均可

在这里插入图片描述

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐