#量化感知训练

在完成 Calibration 之后，如果当前量化结果还不能满足要求，下一步就需要进入量化感知训练。

相比纯浮点模型训练，QAT 由于引入了额外的量化约束，训练过程通常更加困难。量化感知训练工具的目标，是在降低训练复杂度的同时，减少量化模型在部署过程中的工程适配成本。

本阶段的目标，是在已有量化参数初始化的基础上，进一步优化量化误差，为后续模型导出、定点化以及部署验证提供更高精度的 QAT 模型。

#BN 策略与训练模式

开始 QAT 训练前，建议先把 BN 的处理方式和当前训练模式确定下来。这部分配置会直接影响训练行为，也会影响后面看到的精度和收敛表现。

#FuseBN

FuseBN 表示在 Prepare 后、QAT 训练前将 BN 的 weight 和 bias 吸收到 Conv 的 weight 和 bias 中，在训练过程中不再单独更新。这个吸收过程本身是无损的，也是默认的 QAT 训练方式。

适合场景：

• 模型本身就适合融合 BN。

• 您希望训练和最终部署路径尽量保持一致。

#WithBN

WithBN 表示在 QAT 训练阶段保持 Conv + BN 不融合，带着 BN 一起训练，在训练结束后再转换为部署模型时完成融合。

适合场景：

• 浮点训练阶段本身使用了 FreezeBN 或类似策略。

• 您希望 QAT 阶段尽量对齐浮点训练方式。

#相关接口

qat_mode 用于控制当前的量化训练模式。如果您需要切换 BN 策略或训练模式，建议结合当前模型结构、浮点训练策略和后续部署目标一起判断，而不要只凭训练习惯直接套用。

相关接口定义在 horizon_plugin_pytorch.qat_mode 中。结合 qat_mode.py，当前支持三种模式，默认值为 QATMode.FuseBN。

这部分常用到的接口主要有：

• QATMode：用于指定具体的量化训练模式，当前包括 FuseBN、WithBN 和 WithBNReverseFold。

• set_qat_mode：用于设置当前使用的量化训练模式，通常建议在 prepare 之前调用。

• get_qat_mode：用于查看当前已经生效的量化训练模式。

一般可以按下面的思路选择：

• 如果没有明确的特殊需求，通常先从默认的 FuseBN 开始。

• 如果浮点训练阶段使用了 FreezeBN 或类似策略，希望 QAT 阶段尽量保持 BN 行为，可优先尝试 WithBN。

• 如果您明确计划在 QAT 中后期逐步吸收 BN，并希望吸收前后扰动更小，可进一步考虑 WithBNReverseFold。

例如，如果浮点训练阶段使用了 FreezeBN 或类似策略，希望在 QAT 阶段继续保持 BN 参与训练，可以这样设置：

from horizon_plugin_pytorch.qat_mode import QATMode, get_qat_mode, set_qat_mode
from horizon_plugin_pytorch.quantization import prepare

# 默认一般是 FuseBN
print("current qat_mode:", get_qat_mode())

# 在 prepare 之前设置 QAT mode
set_qat_mode(QATMode.WithBN)

# 如果后续计划在 QAT 中后期逐步吸收 BN，
# 也可以改成 QATMode.WithBNReverseFold
# set_qat_mode(QATMode.WithBNReverseFold)

qat_model = prepare(
    float_model,
    example_input,
    qconfig_setter=my_qconfig_setter,
)

如果您还需要继续判断 qat_mode 的选型，或进一步联动调整学习率、scale 更新策略和高精度配置，建议继续阅读 精度调优指导，并重点查看其中的 通用混合精度调优流程、基本调优手段 和 高级调优手段。

#量化参数更新策略

开始训练前，量化参数是否固定、是否持续更新，也会直接影响后面的收敛情况。

当 Calibration 精度较好时，固定 feature map 的量化参数进行 QAT 训练通常可以取得更好的效果。Calibration 精度较差时，则通常不建议直接固定 calibration 得到的量化参数，此时应优先回头排查 Calibration 精度，先解决量化瓶颈。

如果确认 calibration 阶段的量化参数还不能满足精度要求，并且希望在 QAT 训练中持续更新，工具提供了以下更新方式。

#基于统计的方法

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

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

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

#流程和示例

虽然量化感知训练工具并不强制要求您一开始就提供预训练浮点模型，但经验上通常还是建议从预训练的高精度浮点模型开始，这样可以明显降低训练难度。

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())

# 进行量化感知训练
# 作为一个 finetune 过程，量化感知训练一般需要设定较小的学习率
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 往往不是一轮训练就结束，通常需要结合阶段性的训练结果和验证结果，逐步判断当前配置是否合适，以及是否还需要继续训练或调整策略。

因此，除了训练本身，还需要同步做好两件事：

1. 根据当前结果判断是否需要调整学习率、BN 策略、量化参数更新方式或高精度配置。

2. 不只看 QAT model 精度，还要持续验证 quantized model 精度，确认当前训练是否真的有助于后续部署。

从上面的代码可以看到，相比传统的浮点训练，QAT 主要多了两个关键步骤：

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)

到这里，伪量化模型的搭建和参数初始化就完成了，接下来就可以进行常规的训练迭代和模型参数更新，同时持续监控 quantized model 精度。

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

from horizon_plugin_pytorch.quantization import freeze_qat_module

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

#训练过程中重点监控什么

在 QAT 训练过程中，建议至少同时关注下面几类指标：

#QAT model 精度

QAT model 精度主要用于观察训练过程中的整体收敛趋势。它能帮助您判断当前训练是否还在朝更好的方向走，但不能单独作为最终部署精度的判断依据。

#quantized model 精度

quantized model 精度更接近后续部署方向的实际结果，用来判断当前训练结果是否真正有助于后续导出、定点化和部署。

如果只看 QAT model 精度，而不看 quantized model 精度，后续导出或上板阶段更容易出现“训练看起来正常，但部署精度掉点”的情况。

#loss 是否接近 float 训练末期

QAT 本质上通常是对已收敛浮点模型的 finetune，因此 loss 是否与 float 训练末期接近，可以帮助判断量化训练是否进入正常收敛状态。

• loss 过大时，往往意味着难以收敛。

• loss 过小时，也可能带来泛化性问题。

#是否出现 NAN / INF / 长时间不收敛

这类现象通常意味着训练链路本身、scale 更新、量化配置或初始化存在异常，应尽早排查，不建议等训练结束后再统一分析。

#常见异常与建议动作