量化训练流程

本文档仅说明在HAT中进行量化训练时需要的操作，关于量化的基本原理和在训练框架中的实现方式请参阅 horizon_plugin_pytorch 的相关文档。

在量化训练中，由浮点模型到定点模型的转换流程如下：

模型注册，HAT中的各种模块全部采用了注册机制，只有将定义的模型在对应的注册项中进行注册，才可以在config文件中以 dict(type={$class_name}, ...) 的形式使用模型。
qconfig 配置，在对模型量化之前，我们需要为模型配置量化配置（qconfig）， horizon_plugin_pytorch 提供了 QconfigSetter 接口快速配置模型的qconfig, 具体可以参考Qconfig配置.

此外，为使模型可转为量化模型，需要满足一些条件，具体见horizon_plugin_pytorch 的相关文档。

添加自定义模型

import torch
from torch import nn

from hat.registry import OBJECT_REGISTRY


# 使用装饰器的方式将模型进行注册
@OBJECT_REGISTRY.register_module
class ExampleNet(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        ...

添加 config 文件

ckpt_dir = ...

model = dict(type="ExampleNet")

float_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=model,
    data_loader=...,
    optimizer=...,
    batch_processor=...,
    num_epochs=...,
    device=None,
    callbacks=...,
    ...,
)

qat_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode="fuse_bn",
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "float-checkpoint-best.pth.tar"
                ),
            ),
            dict(type="Float2QAT"),
        ],
    ),
    data_loader=...,
    optimizer=...,
    batch_processor=...,
    num_epochs=...,
    device=None,
    callbacks=...,
    ...,
)

训练

只需在使用 tools/train.py 脚本时按顺序指定训练阶段即可，会自动根据训练阶段调用相应的 solver 来执行训练过程：

python3 tools/train.py --stage float ...
python3 tools/train.py --stage qat ...

float：正常的浮点训练。
qat：QAT训练（量化感知训练），首先初始化一个浮点模型，加载训练好的浮点模型权重，再将此模点模型转为QAT模型进行训练, 训练产物是qat伪量化模型。

恢复训练

可以通过在 config 的 {stage}_trainer 中配置 resume_optimizer 和 resume_epoch_or_step 字段来恢复意外中断的训练，或仅恢复optimizer来进行fine-tune。例如：

float_trainer = dict(
    ...
    model_convert_pipeline=dict(
        ...
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path="your_checkpoint_path",  # 要 resume 的 checkpoint 路径
            ),
        ],
    ),
    resume_optimizer=True,      # 恢复 optimizer
    resume_epoch_or_step=True,  # 恢复 epoch 或 step 
    ...,
)

恢复训练有三种使用场景：

完全恢复：该场景为恢复意外中断的训练，会恢复上一个checkpoint的所有状态，包括optimizer、LR、epoch、step 等。该场景只需配置 resume_optimizer 字段即可；
恢复optimizer用于fine-tune：该场景只会恢复optimizer和LR的状态，但epoch、step都会从0开始，用于某些任务的 fine-tune。该场景需要配置 resume_optimizer，并且需要配置resume_epoch_or_step=False。
只加载模型参数：该场景只会加载模型参数，不会恢复其他任何状态(optimizer、epoch、step、LR)。该场景只需要在 model_convert_pipeline 中配置 LoadCheckpoint ，并且需要配置 resume_optimizer=False 和 resume_epoch_or_step=False。

#量化训练流程

#添加自定义模型

#添加 config 文件

#训练

#恢复训练

量化训练流程

添加自定义模型

添加 config 文件

训练

恢复训练