1 前言
在阅读后文plugin使用方式之前，为避免理解歧义，我们先定义一下异构和非异构的概念：异构模型是指部署时一部分运行在 BPU 上，一部分运行在 CPU 上的模型。而非异构模型部署时则完全运行在 BPU 上。
地平线基于PyTorch开发的horizon_plugin_pytorch量化训练工具（该工具将随2023年初的OE开发包释放给XJ3用户）同时支持Eager和fx两种模式。其中，fx模式是从plugin-1.0.0版本之后才开始支持的，相较于Eager方案，他们有如下不同：
eager模式的使用方式建议参考用户手册-4.2量化感知训练章节（4.2.2. 快速上手中有完整的快速上手示例，各使用阶段注意事项建议参考4.2.3. 使用指南）。后文关于异构以及非异构方案的使用方式介绍均基于 fx 模式，fx模式的相关API介绍请参考用户手册-4.2.3.4.2. 主要接口参数说明章节。
2 异构&非异构方案的使用方式
异构以及非异构方案的优缺点如下图所示：
通常来说，我们只会在以下两种情形时使用异构方案：
- 模型中包含 BPU 不支持的算子。
- 模型量化精度误差过大，需要将某些算子放到 CPU 上进行高精度计算。
2.1 非异构方案
由于BPU算子性能远高于CPU算子，从性能优化的角度出发，建议您尽可能使用纯BPU算子搭建模型。
非异构方案的使用步骤大致如下：
1. 浮点模型准备：完成算子替换（参考算子支持列表）、插入量化和反量化节点
2. 设置硬件架构
3. 校准（可选）
4. 模型量化
a. 设置qconfig（推荐先设置全局 qconfig 为get_default_qat_qconfig()，在此基础上根据需求修改，一般而言，只需要对 int16 和高精度输出的 op 单独设置 qconfig）
b. 转qat模型
5. 量化训练&精度验证
6. 转定点模型&精度验证
7. 模型编译
参考代码：
import torch
from horizon_plugin_pytorch.march import March, set_march
from horizon_plugin_pytorch.quantization import (
get_default_calib_qconfig,
get_default_qat_qconfig,
get_default_qat_out_qconfig,
prepare_qat_fx,
convert_fx,
check_model,
compile_model,
)
from torch import nn
from torch.quantization import DeQuantStub, QuantStub
from horizon_plugin_pytorch.utils.onnx_helper import export_to_onnx
# 1.浮点模型准备
class ConvBNReLU(nn.Sequential):
def __init__(self, channels=3):
super(ConvBNReLU, self).__init__(
nn.Conv2d(channels, channels, 1),
nn.BatchNorm2d(channels),
nn.ReLU(),
)
class ConvBNAddReLU(nn.Sequential):
def __init__(self, channels=3):
super(ConvBNAddReLU, self).__init__()
self.conv = nn.Conv2d(channels, channels, 1)
self.bn = nn.BatchNorm2d(channels)
self.relu = nn.ReLU()
self.add = torch.nn.quantized.FloatFunctional()
def forward(self, x):
out = self.conv(x)
out = self.bn(out)
return self.relu(self.add.add(x,out))
class ConvAddReLU(nn.Sequential):
def __init__(self, channels=3):
super(ConvAddReLU, self).__init__()
self.conv = nn.Conv2d(channels, channels, 1)
self.relu = nn.ReLU()
self.add = torch.nn.quantized.FloatFunctional()
def forward(self, x):
out = self.conv(x)
return self.relu(self.add.add(x,out))
class NonHybridModel(nn.Module):
def __init__(self, channels=3):
super().__init__()
# 插入 QuantStub
self.quant = QuantStub()
self.layer0 = ConvBNReLU()
self.layer1 = ConvAddReLU()
self.layer2 = ConvBNAddReLU()
self.conv1 = nn.Conv2d(channels, channels, 1)
# 插入 DequantStub
self.dequant = DeQuantStub()
def forward(self, input):
x = self.quant(input)
x = self.layer0(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.conv1(x)
return self.dequant(x)
data_shape = [1, 3, 224, 224]
data = torch.rand(size=data_shape)
model = NonHybridModel()
# float 模型的推理不要放在 prepare_qat_fx 之后，prepare_qat_fx 会对 float 模型做 inplace 修改
float_res = model(data)
# 2.设置 march
set_march(March.BAYES)
# 3.校准（推荐默认使用）
calibration_model = prepare_qat_fx(
model,
{
# calibration fake quant 只做统计，qat 阶段未使用的 calibration fake quant会被自动去除，不用对高精度输出 op 做高精度设置
"": get_default_calib_qconfig(),
},
)
# calibration 阶段需确保原有模型不会发生变化
calibration_model.eval()
for i in range(5):
calibration_model(torch.rand(size=data_shape))
# 4.模型量化
qat_model = prepare_qat_fx(
calibration_model,
{
"": get_default_qat_qconfig(),
# 设置最后一个conv高精度输出
"module_name": [("conv1", get_default_qat_out_qconfig())]
},
)
# qat 模型的推理不要放在 convert_fx 之后，convert_fx 会对 qat 模型做 inplace 修改
qat_res = qat_model(data)
# 5.量化训练&精度验证
# qat training start
# ... ...
# qat training end
# 6.评测定点精度
quantize_model = convert_fx(qat_model)
quantize_res = quantize_model(data)
# 7.模型编译
# 方式1：将 qat_model 导出为 onnx 之后使用 hb_mapper 工具编译生成 .bin ，算子支持范围对齐 ptq 方案
# 请注意：export_to_onnx 需要在 convert_fx 前使用，convert 会对 qat 模型做 inplace 修改，正确使用时应在第六步之前完成导出，该处仅作代码参考
export_to_onnx(qat_model,data,"qat.onnx",enable_onnx_checker=True, operator_export_type=None)
# 方式2：使用compile_model编译生成.hbm，算子支持范围对齐plugin
traced_model = torch.jit.trace(quantize_model, data)
check_model(quantize_model, data, advice=1)
compile_model(traced_model, [data], opt=3, hbm="./model_output/model.hbm") 2.2 异构方案
异构方案的使用步骤大致如下：
1. 浮点模型准备：
a. 完成算子替换（参考ptq方案算子支持列表）
① 对于非 module 的运算，如果需要单独设置 qconfig 或指定其运行在 CPU 上，需要将其封装成 module，参考后文示例中的_SeluModule
b. 插入量化和反量化节点
① 如果第一个 op 是 cpu op，那么不需要插入 QuantStub
② 如果最后一个 op 是 cpu op，那么可以不用插入 DeQuantStub
2. 设置硬件架构
3. 校准（可选）
4. 模型量化
a. 设置qconfig：推荐先设置全局 qconfig 为get_default_qat_qconfig()，在此基础上根据需求修改，一般而言，只需要对 int16 和高精度输出的 op 单独设置 qconfig
b. 转qat模型：设置 hybrid=True，并通过 hybrid_dict 指定需要运行在cpu上的节点
5. 量化训练&精度验证
6. 导出onnx
7. 评测定点精度
8. 使用hb_mapper工具完成定点转换&模型编译
参考代码：
import numpy as np
import torch
from horizon_plugin_pytorch.march import March, set_march
from horizon_plugin_pytorch.nn import qat
from horizon_plugin_pytorch.quantization import (
get_default_calib_qconfig,
get_default_qat_qconfig,
get_default_qat_out_qconfig,
prepare_qat_fx,
convert_fx,
)
from torch import nn
from torch.quantization import DeQuantStub, QuantStub
from horizon_plugin_pytorch.utils.onnx_helper import export_to_onnx
# 1.浮点模型准备
class ConvReLU(nn.Sequential):
def __init__(self, channels=3):
super().__init__()
self.conv = nn.Conv2d(channels, channels, 1)
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.conv(x)
x = self.relu(x)
return x
# 封装 functional selu 为 module，便于单独设置
class _SeluModule(nn.Module):
def forward(self, input):
return torch.nn.functional.selu(input)
class HybridModel(nn.Module):
def __init__(self, channels=3):
super().__init__()
# 第一层为cpu节点
#self.quant = QuantStub()
self.layer0 = ConvBNReLU()
self.layer1 = ConvBNReLU()
self.layer2 = ConvBNReLU()
self.selu = _SeluModule()
self.conv0 = nn.Conv2d(channels, channels, 1)
self.conv1 = nn.Conv2d(channels, channels, 1)
# 插入 DequantStub
self.dequant = DeQuantStub()
def forward(self, input):
x = self.selu(x)
x = self.layer0(x)
x = self.selu(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.conv0(x)
x = self.conv1(x)
x = self.selu(x)
return self.dequant(x)
data_shape = [1, 3, 224, 224]
data = torch.rand(size=data_shape)
model = HybridModel()
# float 模型的推理不要放在 prepare_qat_fx 之后，prepare_qat_fx 会对 float 模型做 inplace 修改
float_res = model(data)
# 2.设置 march
set_march(March.BAYES)
# 3.校准（推荐默认使用）
calibration_model = prepare_qat_fx(
model,
{
# calibration fake quant 只做统计，qat 阶段未使用的 calibration fake quant会被自动去除，不用对高精度输出 op 做高精度设置
"": get_default_calib_qconfig(),
},
hybrid=True,
hybrid_dict={
"module_name": ["layer1.conv", "conv0"],
"module_type": [_SeluModule],
},
)
# calibration 阶段需确保原有模型不会发生变化
calibration_model.eval()
for i in range(5):
calibration_model(torch.rand(size=data_shape))
# 4.模型量化
qat_model = prepare_qat_fx(
calibration_model,
{
"": get_default_qat_qconfig(),
# 设置最后一个conv高精度输出
"module_name": [("conv1", get_default_qat_out_qconfig())]
},
hybrid=True,
hybrid_dict={
"module_name": ["layer1.conv", "conv0"],
"module_type": [_SeluModule],
},
)
# qat 模型的推理不要放在 convert_fx 之后，convert_fx 会对 qat 模型做 inplace 修改
qat_res = qat_model(data)
# 5.量化训练&精度验证
# qat training start
# ... ...
# qat training end
# 6.导出onnx
export_to_onnx(qat_model,data,"qat.onnx",enable_onnx_checker=True, operator_export_type=None)
# 7.评测定点精度
# 异构方案的quantize_model仅可用于评测精度，无法用于模型编译；也可使用hb_mapper生成的quantized.onnx评测精度
quantize_model = convert_fx(qat_model)
quantize_res = quantize_model(data) 8. 使用hb_mapper工具完成定点转换&模型编译
最简config.yaml配置示例：
model_parameters:
onnx_model: "qat.onnx"
march: "bayes"
output_model_file_prefix: 'hybrid'
input_parameters:
input_type_rt: 'featuremap'
input_type_train: 'featuremap'
input_layout_train: 'NCHW'
input_layout_rt: 'NCHW'
norm_type: 'no_preprocess'
calibration_parameters:
calibration_type: 'load'
# conv0（导出onnx后name为Conv_32）是bpu可支持的算子，虽然它自身没有FakeQuantize，但其输入为量化输入，转换工具仍然会尝试量化该节点，因此需要手动将其指定到cpu上。而layer1.conv的输入为浮点输入，因此无需手动指定。
run_on_cpu: "Conv_32"
compiler_parameters:
compile_mode: 'latency'
debug: false
optimize_level: 'O3' 模型转换：
hb_mapper makertbin -c config.yaml --model-type onnx 从转换日志可见，selu、layer1.conv以及conv0都运行到了cpu上：
使用hb_perf工具生成模型结构图，可观察到第一段以及第二段bpu结尾的conv未设置高精度输出，仅对第三段bpu的尾部节点设置了高精度输出：
3 其他常见问题
- 打印qat_model发现多出了一些generated_add节点，这是为什么？
答：这是因为模型中直接使用了“+”号运算符，工具会通过新注册的generated_add来实现将其自动替换为FloatFunctional.add。建议大家最好在浮点模型准备阶段自行完成算子替换，因为工具自动转换时会依据执行顺序为add命名，若您后续修改了模型，可能会出现无法加载原ckpt的问题。