QAT快速上手(eager mode)

目录

  • 1. 必要步骤速览-

  • 2. 各步骤详解-

    • 2.1 浮点模型准备-

    • 2.2 数据校准-

    • 2.3 量化训练-

    • 2.4 定点转换-

    • 2.5 模型编译-

  • 3. 常见问题-


Pytorch官方提供了两种量化模式:Eager Mode Quantization 和 FX Graph Mode Quantization,这两种模式的区别如下(来源:PyTorch官方文档):-

-
对于新手而言,还是优先推荐 FX mode(QAT快速上手(fx mode)),虽然 FX mode 也有一些限制,需要模型满足 “symbolically traceable” 的要求,但整体而言其自动化程度会高很多,无需用户手写算子融合、手动替换不支持的算子,若碰到无法解决的问题,再回退到 Eager mode。

1. 必要步骤速览

整个QAT方案从浮点到部署模型就包括五个步骤:浮点模型准备、数据校准、量化训练(可选)、定点转换、模型编译。必要的步骤和示例代码如下所示,每个步骤的详细说明和注意事项可参考后文讲解。完整的示例推荐参考OE开发包中/ddk/samples/ai_toolchain/horizon_model_train_sample/plugin_basic目录下的eager_mode.py脚本。-

强烈建议在量化训练前(甚至是浮点模型设计阶段)先跳过训练过程完整走完prepare->convert->check步骤,确保模型可被硬件支持。

from horizon_plugin_pytorch.quantization import (
    convert, 
    prepare_qat,
    set_fake_quantize,
    FakeQuantState,
    check_model,
    compile_model,
)
from horizon_plugin_pytorch.quantization.qconfig import (
    default_calib_8bit_fake_quant_qconfig,
    default_qat_8bit_fake_quant_qconfig,
    default_calib_8bit_weight_32bit_out_fake_quant_qconfig,
    default_qat_8bit_weight_32bit_out_fake_quant_qconfig
)
from horizon_plugin_pytorch.march import March, set_march
import torch

set_march(March.BAYES) # 在prepare之前设置计算架构
# 1.准备浮点模型(对浮点模型需要进行一些必要的改造,具体请见后文 2.1 节)
float_model = load_float_model(pretrain=True) 
float_model.fuse_model()
ori_float_model = float_model
# 2.数据校准
float_model = copy.deepcopy(ori_float_model)
float_model.qconfig = default_calib_8bit_fake_quant_qconfig
# 配置高精度输出,具体原因请见后文常见问题1
# plugin ≤ v1.6.2 配置高精度输出需使用 default_calib_out_8bit_fake_quant_qconfig
float_model.classifier.qconfig = default_calib_8bit_weight_32bit_out_fake_quant_qconfig
calib_model = prepare_qat(float_model)
calib_model.eval()
# 关于FakeQuantState相关说明请参考后文第三章常见问题 3
set_fake_quantize(calib_model, FakeQuantState.CALIBRATION)
calibrate(calib_model)
# 评测数据校准精度
calib_model.eval()
set_fake_quantize(calib_model, FakeQuantState.VALIDATION)
evaluate(calib_model)
torch.save(calib_model.state_dict(), "calib-checkpoint.ckpt")
# 3.量化训练(若数据校准精度已达标,可跳过该步骤)
float_model = copy.deepcopy(ori_float_model)
float_model.qconfig = default_qat_8bit_fake_quant_qconfig
float_model.classifier.qconfig = default_qat_8bit_weight_32bit_out_fake_quant_qconfig
qat_model = prepare_qat(float_model) 
qat_model.load_state_dict(calib_model.state_dict())
qat_model.train()
set_fake_quantize(qat_model, FakeQuantState.QAT)
train(qat_model)
# 评测量化训练精度
qat_model.eval()
set_fake_quantize(qat_model, FakeQuantState.VALIDATION)
evaluate(qat_model)
# 4.定点转换
base_model = qat_model # 校准精度满足预期时使用 calib_model
quantized_model = convert(base_model)
# 评测定点精度(正常情况下该模型精度与hbm模型精度一致,因此请以此来评价最终部署精度)
evaluate(quantized_model)
# 5.模型编译(注意模型和数据要放在 CPU 上)
script_model = torch.jit.trace(quantized_model.cpu(), example_input)
check_model(script_model, [example_input])
compile_model(script_model,[example_input],hbm="model.hbm",input_source="pyramid",opt=3)

2. 各步骤详解

2.1 浮点模型准备

**a. 请使用足够的数据量将浮点模型正常训练至收敛后再进行量化训练。-

b. 强烈建议对输入数据进行归一化处理,有利于浮点收敛的同时也可使得模型对量化更友好。-

c. 建议您在浮点模型设计阶段对照算子支持列表,避免使用不支持的算子导致后续prepare qat或者编译报错。-

d. 更多关于如何搭建量化友好模型的说明可参考用户手册 4.2.4.1浮点模型的要求**

若选择eager mode,那在浮点模型准备阶段需要完成下面几个必要步骤(可参考Pytorch官方文档):

  • 在模型输入前插入 QuantStub节点,在模型输出后插入 DequantStub节点。有如下注意事项:

    • 多个输入仅在 scale 相同时可以共享 QuantStub,否则请为每个输入定义单独的 QuantStub
    • 建议使用horizon_plugin_pytorch.quantization.QuantStub 默认动态统计输入scale,若是可提前计算出scale的场景建议手动设置scale(例如bev模型的homo矩阵),对应的公版接口torch.quantization.QuantStub不支持手动设置。
  • 参考算子约束列表完成算子替换。做这一步的原因是eager mode对算子的支持有一定限制,对部分算子的量化会有一些特殊处理。(例如需要用torch.nn.ReLU替换torch.nn.functional.relu,add/cat/matmul等需要替换为horizon.nn.quantized.FloatFunctional,具体可参考如下示例)。

    from horizon_plugin_pytorch.nn.quantized import FloatFunctional

    class ModelName(nn.Module):
    def init(···):
    self.matmul1 = FloatFunctional()
    def forward(···):
    result = self.matmul1.matmul(a, b)
    ···

  • 算子融合:算子融合可显著提高模型量化精度及部署性能,因此这是不可或缺的步骤,目前可支持conv/ConvTranspose2d/linear、bn、relu/relu6、add之间的融合(具体可参考:PythonPath/horizon_plugin_pytorch/quantizationfuse_modules.py),支持使用算子下标或算子名称来写融合函数,具体可参考如下示例。

    from horizon_plugin_pytorch.quantization import fuse_known_modules

    class ConvBnRelu(nn.Module):
    def init(self):
    super().init()
    self.conv = nn.Conv2d(2, 2, 1, bias=None)
    self.bn = nn.BatchNorm2d(2)
    self.relu = nn.ReLU(inplace=True)
    def forward(self, x):
    ···
    def fuse_model(self):
    # 方式1:使用算子下标
    torch.quantization.fuse_modules(
    self,
    [“0”, “1”, “2”],
    inplace=True,
    fuser_func=fuse_known_modules
    )
    # 方式2:使用算子名称
    torch.quantization.fuse_modules(
    self,
    [“conv”, “bn”, “relu”],
    inplace=True,
    fuser_func=fuse_known_modules,
    )

2.2 数据校准

对于部分模型,仅通过 Calibration 便可使精度达到要求,不必进行比较耗时的量化感知训练。即使模型经过量化校准后无法满足精度要求,此过程也可降低后续量化感知训练的难度,缩短训练时间,提升最终的训练精度。数据校准的具体配置方式及调参建议可参考QAT方案Calibration使用说明

2.3 量化训练

量化训练一些推荐的超参配置如下表所示:

超参

推荐配置

高级配置(如果推荐配置无效请尝试)

LR

从0.001开始,搭配StepLR做2次scale=0.1的lr decay

1. 调整lr在0.0001->0.001之间,配合1-2的lr decay。-
2. LR 更新策略也可以尝试把 StepLR 替换为 CosLR。-
3. QAT使用AMP,适当调小lr,过大导致nan。

Epoch

浮点epoch的10%

1. 根据loss和metric的收敛情况,考虑是否需要适当延长epoch。

Weight decay

与浮点一致

1. 建议在4e-5附近做适当调整。weight decay过小导致weight方差过大,过大导致输出较大的任务输出层weight方差过大。

optimizer

与浮点一致

1. 如果浮点训练采用的是 OneCycle 等会影响 LR 设置的优化器,建议不要与浮点保持一致,使用 SGD 替换。

transforms(数据增强)

与浮点一致

1. QAT阶段可以适当减弱,比如分类的颜色转换可以去掉,RandomResizeCrop的比例范围可以适当缩小

averaging_constant(qconfig_params)

1. 使用calibration后推荐减弱激活更新:-
weight averaging_constant=1.0-
activation averaging_constant=0.0

1. calibration的精度和浮点差距较大时:activation averaging_constant不要设置成0.0-
2. weight averaging_constant一般不需要设置成0.0,实际情况可以在(0,1.0]之间调整

强烈建议您先尝试数据校准,若精度不满足预期再进行量化训练(注意要加载数据校准后权重参数)。-
量化训练阶段的调参建议可以参考用户手册 量化训练精度调优建议

2.4 定点转换

请注意,定点模型和伪量化模型之间无法做到完全数值一致,所以请以定点模型的精度为准。若定点精度不达标,仍需要继续进行量化训练,建议多保留几个epoch的qat模型权重,便于寻找最优的定点精度。(qat或者calibrate精度高并不一定代表定点精度高,可以考虑进行一些回退,平衡最终的定点精度)

在正常情况下,定点模型的精度与板端部署精度是可以保持完全一致的,因此可使用该模型来评测最终部署精度。

2.5 模型编译

模型编译阶段包括以下三个步骤:

script_model = torch.jit.trace(quantized_model, example_input)
check_model(script_model.cpu(), [example_input])
compile_model(script_model,[example_input],hbm="model.hbm",input_source="pyramid",opt=O3)

compile_model()更多配置项请参考 用户手册-模型编译

其中trace之后生成的script_model可以使用horizon_plugin_pytorch.jit.save接口保存后迁移到其他机器上进行推理评测,由于推理保存后的模型要求device和trace时保持一致,使用to(device)操作修改可能会出现forward报错,具体原因和推荐的解决方式可参考用户手册-量化部署 PT 模型的跨设备 Inference 说明。-
若使用了rgb/bgr格式训练模型,部署时设置input_souce为pyramid或resizer,需要在trace之前手动插入预处理节点centered_yuv2rgb和centered_yuv2bgr ,具体可参考用户手册-RGB888 数据部署

3. 常见问题

1. 为何要设置高精度输出?-
答:依据神经网络量化背景中的介绍可知乘法累加器计算得到的激活值是int32的,为了让下一层op可以继续计算,会经过requantization的操作转为int8/int16,因此若最后一层是conv/linear节点的话,建议设置高精度输出,使得模型可以以int32格式直接输出,对精度保持情况大有裨益。

plugin ≤ v1.6.2 配置高精度输出需使用 default_calib_out_8bit_fake_quant_qconfig ,但该参数将在后续版本中被废弃

2. 如何理解fake quantize的几种状态?-
fake quantize 一共有三种状态,分别需要在 QATcalibrationvalidation 前使用set_fake_quantize将模型的 fake quantize 设置为对应的状态。在 calibration 状态下,仅观测各算子输入输出的统计量。在 QAT 状态下,除观测统计量外还会进行伪量化操作。而在 validation 状态下,不会观测统计量,仅进行伪量化操作。-
感兴趣的话大家可以到python path下找到horizon_plugin_pytorch/quantization/fake_quantize.py,查看set_fake_quantize的实现,以下为截取的关键片段:

···
if mode == FakeQuantState.QAT:
    assert (
        mod.training
    ), "Call model.train() before set fake quant to QAT mode."
    enable_fake_quant(mod)
elif mode == FakeQuantState.CALIBRATION:
    assert (
        not mod.training
    ), "Call model.eval() before set fake quant to CALIBRATION mode."
    disable_fake_quant(mod)
    if isinstance(
        mod, FakeQuantizeBase
    ) or _is_fake_quant_script_module(mod):
        mod.train()
elif mode == FakeQuantState.VALIDATION:
    assert (
        not mod.training
    ), "Call model.eval() before set fake quant to VALIDATION mode."
    # observer won't work in eval mode
    enable_fake_quant(mod)

从这段代码中我们可以得知:

  • 在QAT之前要求模型设置的是train()的状态,CALIBRATION和VALIDATION则要求模型是eval()状态,这主要是为了使bn、dropout等处于正确的状态(训练的时候bn会更新,评测的时候bn不更新)。
  • CALIBRATION时会disable_fake_quant,并设置fake_quant状态为train(),即不进行伪量化操作,仅观测算子输入输出统计量,更新scale;QAT时会观测统计量并进行伪量化操作;VALIDATION时不会观测统计量,仅进行伪量化操作。

因此如下常见误操作会导致一些异常现象:

  • 数据校准之前模型设置为train()的状态,且未使用set_fake_quantize,等于是在跑QAT训练;
  • 数据校准之前模型设置为eval()的状态,且未使用set_fake_quantize,会导致scale一直处于初始状态,全为1;
  • 数据校准之前模型设置为eval()的状态,且正确使用了set_fake_quantize,但是在这之后又设置了一遍model.eval(),这将导致fake_quant未处于训练状态,scale一直处于初始状态,全为1;

3. 如何固定模型某个部分参数不更新?-
答:特别是在多任务模型场景,新增head后再次对模型进行量化,不想改动backbone和已有head的量化参数,依据上一题的介绍,可在set_fake_quantize(model, FakeQuantState.CALIBRATION)或者set_fake_quantize(model, FakeQuantState.QAT)之后再执行一次model.backbone.eval()即可,则backbone的量化参数将不会更新。

1 个赞