量化训练流程

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

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

其中大部分步骤都已集成在HAT的训练pipeline中，只需注意在添加自定义模型时实现 fuse_model 方法来完成模型融合，且实现 set_qconfig 方法对量化方式进行配置即可。在编写模型时需要注意以下几点：

HAT只会调用最外层模块的 fuse_model 方法，因此在 fuse_model 的实现中要负责所有子模块的fuse。
优先使用 hat.models.base_modules 中提供的基础模块，这些基础模块已实现 fuse_model 方法，可减少工作量和开发难度。
模型注册，HAT中的各种模块全部采用了注册机制，只有将定义的模型在对应的注册项中进行注册，才可以在config文件中以 dict(type={$class_name}, ...) 的形式使用模型。
需要在最外层模块实现 set_qconfig 方法，如果子模块中有特殊layer需单独设置 QConfig，也需要在该子模块中实现 set_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):
        ...

    def fuse_model(self):
        # 需要调用所有子模块的 fuse_model 方法
        if hasattr(self.submodule, "fuse_model"):
            self.submodule.fuse_model()

        # 具体 fuse 的接口见 horizon_plugin_pytorch 文档
        ...

    def set_qconfig(self):
        # 具体模型量化配置的接口见 horizon_plugin_pytorch 文档
        from hat.utils import qconfig_manager
        # 默认使用 qconfig_manager.get_default_qat_qconfig() 得到的 QConfig
        self.qconfig = qconfig_manager.get_default_qat_qconfig()
        # 对需要特殊处理的子模块，调用子模块的 set_qconfig，
        # 子模块的 set_qconfig  中只需实现对特殊 layer 的 QConfig 设置
        if hasattr(self.submodule, "set_qconfig"):
            self.submodule.set_qconfig()
        # 如果有特殊节点不需要设置QConfig，比如 loss，需要设置其QConfig 为 None
        if self.loss is Not None:
            self.loss.qconfig = None
        ...

添加 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=...,
    ...,
)

val_callback = dict(
    type="Validation",
    data_loader=...,
    batch_processor=...,
    callbacks=[val_metric_updater, ...],
)

trace_callback = dict(
    type="SaveTraced",
    save_dir=ckpt_dir,
    trace_inputs=deploy_inputs,
)

ckpt_callback = dict(
    type="Checkpoint",
    save_dir=ckpt_dir,
    name_prefix=training_step + "-",
    strict_match=True,
    mode="max",
)

int_trainer = dict(
    type="Trainer",
    model=deploy_model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode="fuse_bn",
        converters=[
            dict(type="Float2QAT"),
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "qat-checkpoint-best.pth.tar"
                ),
            ),
            dict(type="QAT2Quantize"),
        ],
    ),
    # int_trainer 中实际不包含训练流程
    data_loader=None,
    optimizer=None,
    batch_processor=None,
    num_epochs=0,
    ################################
    device=None,
    callbacks=[
        ckpt_callback,
        trace_callback,
    ],
)

训练

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

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

float：正常的浮点训练。
qat：QAT训练（量化感知训练），首先初始化一个浮点模型，加载训练好的浮点模型权重，再将此模点模型转为QAT模型进行训练。
int_infer：定点转化预测，此阶段首先初始化一个浮点模型，将此浮点模型先转为QAT模型并加载训练好的QAT模型权重，再将 QAT模型转为定点模型。转出的定点模型无法进行训练，只能执行validation得到最终的定点模型精度。

恢复训练

可以通过在 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 文件

训练

恢复训练