PyTorch 量化

一、量化是什么？

量化是一种将深度学习模型中的浮点数运算转换为整数运算的技术。在量化过程中，模型的权重和激活值从 32 位浮点数（FP32）转换为 8 位整数（INT8），从而将模型大小减少约 4 倍，并显著降低内存带宽需求。这使得量化模型在推理阶段运行速度更快，尤其是在支持高性能矢量化操作的硬件平台上。

二、PyTorch 量化的优势

1. 模型体积更小：INT8 量化可使模型文件大小减小至原来的 1/4，便于在设备端部署，如移动设备和嵌入式设备。
2. 推理速度快：量化模型在推理时，因减少了内存访问和利用了整数运算的硬件加速，通常比 FP32 模型快 2 到 4 倍。
3. 能耗降低：整数运算相比浮点运算能耗更低，适合在资源受限的环境中运行。

三、PyTorch 量化的三种主要方法

（一）训练后动态量化

训练后动态量化主要对模型的权重进行量化，激活值在推理过程中动态量化。这种方法适合小批量的 LSTM 和 Transformer 类型模型，因为这类模型的执行时间主要由从内存中加载权重决定。

示例代码：

import torch
import torch.quantization


## 创建一个简单的模型
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = torch.nn.Linear(10, 10)


    def forward(self, x):
        return self.linear(x)


model = SimpleModel()
## 将模型转为评估模式
model.eval()


## 动态量化模型
quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)


## 测试量化模型
input_tensor = torch.randn(1, 10)
output = quantized_model(input_tensor)
print(output)

（二）训练后静态量化

训练后静态量化不仅对权重进行量化，还基于校准数据集预先计算激活张量的比例因子和偏差。这种方法适用于 CNN 等模型，可充分利用内存带宽和计算节省。

步骤如下：

1. 准备模型：添加 QuantStub 和 DeQuantStub 模块，指定量化激活和反量化的位置；确保不重复使用模块；将需要重新量化的操作转换为模块形式。
2. 模块融合：将操作组合（如 Conv+ReLU）融合为一个模块，提高准确性和性能。
3. 指定量化配置：选择量化方法，如对称或非对称量化，以及校准技术。
4. 校准模型：在准备好的模型上运行校准数据集。
5. 转换模型：将校准后的模型转换为量化模型。

示例代码：

## 继续使用上面定义的 SimpleModel
model = SimpleModel()
model.eval()


## 添加量化和反量化存根
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)


## 假设有一个校准数据集 calibrate_data
calibrate_data = [torch.randn(1, 10) for _ in range(10)]
with torch.no_grad():
    for data in calibrate_data:
        model(data)


## 转换模型为量化模型
torch.quantization.convert(model, inplace=True)


## 测试量化模型
output = model(input_tensor)
print(output)

（三）量化意识训练（QAT）

量化意识训练在训练过程中模拟量化效果，通过伪量化模块对量化误差进行建模，使模型在训练时就适应量化带来的变化。

步骤如下：

1. 准备模型：与训练后静态量化类似，添加量化和反量化存根，指定量化配置。
2. 融合模块：融合操作组合。
3. 插入伪量化模块：使用 torch.quantization.prepare_qat() 插入伪量化模块。
4. 训练模型：使用常规训练方法训练模型，伪量化模块会模拟量化效果。
5. 转换模型：训练完成后，将模型转换为量化模型。

示例代码：

model = SimpleModel()
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)


## 训练模型（示例中省略训练循环）


## 转换模型为量化模型
torch.quantization.convert(model, inplace=True)


## 测试量化模型
output = model(input_tensor)
print(output)

四、量化张量操作

PyTorch 提供了多种量化张量操作，用于在量化模型中进行数据处理。以下是一些常用的操作：

（一）量化和反量化

1. torch.quantize_per_tensor() ：使用每个张量的标度和零点将浮点张量转换为量化张量。
2. torch.quantize_per_channel() ：使用每个通道的标度和零点将浮点张量转换为量化张量。
3. torch.dequantize() ：将量化张量转换为浮点张量。

示例代码：

float_tensor = torch.randn(2, 3)
scale = 1.0
zero_point = 0
quantized_tensor = torch.quantize_per_tensor(float_tensor, scale, zero_point, torch.qint8)
dequantized_tensor = torch.dequantize(quantized_tensor)

（二）基本操作

1. torch.ops.quantized.add() ：对两个量化张量进行加法运算。
2. torch.ops.quantized.cat() ：将多个量化张量沿指定维度连接。
3. torch.ops.quantized.mul() ：对两个量化张量进行乘法运算。

示例代码：

quantized_tensor1 = torch.quantize_per_tensor(torch.randn(2, 3), 1.0, 0, torch.qint8)
quantized_tensor2 = torch.quantize_per_tensor(torch.randn(2, 3), 1.0, 0, torch.qint8)
result_add = torch.ops.quantized.add(quantized_tensor1, quantized_tensor2)
result_mul = torch.ops.quantized.mul(quantized_tensor1, quantized_tensor2)

五、量化模型的部署

量化模型可以部署在各种支持量化运算的硬件平台上，如具有 AVX2 支持的 x86 CPU 和 ARM CPU。在部署之前，确保量化配置与目标硬件平台匹配，以实现最佳性能。

六、总结

通过本教程，我们深入探讨了 PyTorch 中的量化技术，包括三种主要的量化方法、量化张量操作以及量化模型的部署。量化技术在模型压缩和加速推理方面具有显著优势，尤其适用于设备端部署。选择合适的量化方法并正确应用，可以有效提高模型的性能和效率。

如果你希望进一步学习 PyTorch 量化技术或深入了解其他深度学习相关内容，可以访问编程狮（w3cschool.cn）网站，那里提供了丰富的教程和资源。