一、定义
- 分类
- 模型量化接口进阶
- 量化支持的模式/流/硬件
二、实现
- 分类
- 模型量化接口进阶
https://pytorch.org/docs/stable/quantization-support.html
2.1 算子融合
定义: 将多个算子融合到一起,运算时可以加快运行速度。
import torch
# define a floating point model where some layers could benefit from QAT
class M(torch.nn.Module):
def __init__(self):
super().__init__()
# QuantStub converts tensors from floating point to quantized
self.quant = torch.ao.quantization.QuantStub()
self.conv = torch.nn.Conv2d(1, 1, 1)
self.bn = torch.nn.BatchNorm2d(1)
self.relu = torch.nn.ReLU()
# DeQuantStub converts tensors from quantized to floating point
self.dequant = torch.ao.quantization.DeQuantStub()
def forward(self, x):
x = self.quant(x)
x = self.conv(x)
x = self.bn(x)
x = self.relu(x)
x = self.dequant(x)
return x
model_fp32 = M()
# model must be set to eval for fusion to work
model_fp32.eval()
print(model_fp32)
model_fp32_fused = torch.ao.quantization.fuse_modules(model_fp32,
[['conv', 'bn', 'relu']])
print(model_fp32_fused)
3. 量化支持的模式/流/硬件
选择条件: 模式支持、硬件支持、算子支持等。
3.1 量化模式选择
3.2 硬件模式选择
3.3 算子对模式的支持