• 中文
  • 3.8.1

    • 量化感知训练

    量化感知训练通过在模型中插入一些伪量化节点，从而使得通过量化感知训练得到的模型转换成定点模型时尽可能减少精度损失。 量化感知训练和传统的模型训练无异，可以从零开始，搭建一个伪量化模型，然后对该伪量化模型进行训练。 由于部署的硬件平台有诸多限制，搞清这些限制，并且根据这些限制搭建伪量化模型门槛较高。 量化感知训练工具通过在您提供的浮点模型上根据部署平台的限制自动插入伪量化量化算子的方法，降低开发量化模型的门槛。

    量化感知训练由于施加了各种限制，因此，一般来说，量化感知训练比纯浮点模型的训练更加困难。量化感知训练工具的目标是降低量化感知训练的难度，降低量化模型部署的工程难度。

    注解

    多机多卡的适配

    calibration 和 量化感知训练的过程均支持多机多卡，任务启动方式和浮点模型完全相同，每个进程中在将模型使用 torch.nn.parallel.DistributedDataParallel 封装之前先进行 prepare 即可。

    流程和示例

    虽然我们的量化感知训练工具并不强制要求您在最开始提供预训练的浮点模型，但是，经验表明，通常从预训练的高精度浮点模型开始量化感知训练能大大降低训练的难度。

    from horizon_plugin_pytorch.quantization import FakeQuantState, prepare, set_fake_quantize
from horizon_plugin_pytorch.quantization import hbdk4 as hb4
from horizon_plugin_pytorch.quantization.qconfig_setter import *
from hbdk4.compiler import convert

# 将模型转为 QAT 状态
qat_model = prepare(
    float_model,
    example_input,
    qconfig_setter=QconfigSetter(
        get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
        ],
    ),
).to(device)
# 加载 Calibration 模型中的量化参数
qat_model.load_state_dict(calib_model.state_dict())
# 进行量化感知训练
# 作为一个 filetune 过程，量化感知训练一般需要设定较小的学习率
optimizer = torch.optim.SGD(
    qat_model.parameters(), lr=0.0001, weight_decay=2e-4
)

for nepoch in range(epoch_num):
    # 注意此处对 QAT 模型 training 状态的控制方法
    qat_model.train()
    set_fake_quantize(qat_model, FakeQuantState.QAT)

    train_one_epoch(
        qat_model,
        nn.CrossEntropyLoss(),
        optimizer,
        None,
        train_data_loader,
        device,
    )

    # 注意此处对 QAT 模型 eval 状态的控制方法
    qat_model.eval()
    set_fake_quantize(qat_model, FakeQuantState.VALIDATION)

    # 测试 qat 模型精度
    top1, top5 = evaluate(
        qat_model,
        eval_data_loader,
        device,
    )
    print(
        "QAT model: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
            top1.avg, top5.avg
        )
    )

# 测试 quantized 模型精度
qat_hbir_model = hb4.export(
    qat_model. example_input
)
quantized_hbir_model = convert(qat_hbir_model)

top1, top5 = evaluate(
    quantized_hbir_model,
    eval_data_loader,
)
print(
    "Quantized model: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
        top1.avg, top5.avg
    )
)
    注意

    由于部署平台的底层限制，QAT 模型无法完全代表最终上板精度，请务必监控 quantized 模型精度，确保 quantized 模型精度正常，否则可能出现模型上板掉点问题。

    由上述示例代码可以看到，与传统的纯浮点模型训练相比，量化感知训练多了两个步骤：

    1. prepare。这一步骤的目标是对浮点网络进行变换，插入伪量化节点。

    2. 加载 Calibration 模型参数。通过加载 Calibration 得到的伪量化参数，来获得一个较好的初始化。

      对于 state_dict 的修改不应当只关注 key 和 value，同样需要考虑 _metadata。如下是一个复制 state_dict 的例子：

      new_state_dict = OrderedDict()
for k, v in state_dict.items():
    new_state_dict[k] = v

if hasattr(state_dict, "_metadata"):
    new_state_dict._metadata = copy.deepcopy(state_dict._metadata)
    注意

    算子的兼容性依赖 torch.nn.Module 的 _version 变量实现，_version 会在 state_dict._metadata 中保存，请确保在保存或加载 state_dict 的过程中保留了 _metadata，否则可能引起兼容性问题。

    至此，已经完成了伪量化模型的搭建和参数的初始化，接下来就可以进行常规的训练迭代和模型参数更新，并且监控 quantized 模型精度。

    为了满足分段部署的需要或对齐浮点训练策略，训练过程中可能需要固定模型的某些部分，可以参考下面的做法进行固定。

    from horizon_plugin_pytorch.quantization import freeze_qat_module

# 模型权重 / 量化参数会被固定，所有算子被设置为 eval 状态。
# 需要保证调用 freeze_qat_module 接口在调用 train() / eval() / set_fake_quantize() 等会改变模型状态的接口之后。
freeze_qat_module(model)

    量化参数更新策略

    Calibration 精度较好时，固定 feature map 的量化参数进行 QAT 训练可以取得更好的效果。精度较差时，则通常不能固定 calibration 得到的量化参数（此时建议优先调试 calibration 精度，解决量化瓶颈）。若确认 calibraiton 阶段的量化参数无法满足精度要求，希望在 qat 训练中持续更新的话，工具提供了以下更新方式。

    基于统计的方法

    出于性能考虑，我们推荐在 QAT 阶段使用 MinMaxObserver，对量化参数进行持续统计和更新。使用方法请参考 QConfig 详解。

    基于学习的方法（不推荐）

    工具支持 LSQ 算法（详细信息参见 Learned Step Size Quantization），对应的接口为 _LearnableFakeQuantize，使用方法请参考 FakeQuantize 的定义。