• English
  • 3.5.0

    • Quantized Training Process Introduction

    This document only describes the operations needed to perform quantization training in HAT. For the basic principles of quantization and its implementation in the training framework, refer to the documentation of horizon_plugin_pytorch .

    In quantization training, the conversion process from a floating-point model to a fixed-point model is as follows:

    qat

    Most of these steps are already integrated in the HAT training pipeline, and you only needs to pay attention to implementing the fuse_model method to complete the model fusion when adding a custom model and the set_qconfig method to configure the quantization method. The following points need to be noted when writing the models.

    • HAT will only call the fuse_model method of the outermost module, so the implementation of fuse_model is responsible for the fuse of all submodules.

    • Preference should be given to the base modules provided in hat.models.base_modules, which has already implemented the fuse_model method to reduce the effort and development difficulties.

    • Model registration, all the modules in HAT use the registration mechanism, only when the defined model is registered in the corresponding registration item, can the model be used in the config file as dict(type={$class_name}, ...) .

    • The set_qconfig method needs to be implemented in the outermost module. If there is a special layer in a submodule that needs a separate QConfig setting, the set_qconfig method needs to be implemented in that submodule as well, details of which can be found in the Writing Specifications of set_qconfig and Customization of qconfig sections.

    In addition, to make the model transferable to a quantized model, some conditions need to be met, as described in the documentation for horizon_plugin_pytorch.

    Add Custom Models

    import torch
from torch import nn

from hat.registry import OBJECT_REGISTRY


# Registering the model using decorators
@OBJECT_REGISTRY.register_module
class ExampleNet(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        ...

    def fuse_model(self):
        # You need to call the fuse_model method of all the submodules
        if hasattr(self.submodule, "fuse_model"):
            self.submodule.fuse_model()

        # Descriptions on the fuse interface can be found in the horizon_plugin_pytorch document
        ...

    def set_qconfig(self):
        # Descriptions on the interface of the specific model quantization configuration can be found in the horizon_plugin_pytorch document
        from hat.utils import qconfig_manager
        # The default QConfig is obtained by using qconfig_manager.get_default_qat_qconfig()
        self.qconfig = qconfig_manager.get_default_qat_qconfig()
        # For submodules requiring special handling, call the submodule's set_qconfig
        # For set_qconfig of the submodules, you only need to implement the QConfig settings for the special layer
        if hasattr(self.submodule, "set_qconfig"):
            self.submodule.set_qconfig()
        # If a special node does not require any QConfig setting, such as loss, you need to set its QConfig to None
        if self.loss is Not None:
            self.loss.qconfig = None
        ...

    Add the Config File

    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 contains no training process
    data_loader=None,
    optimizer=None,
    batch_processor=None,
    num_epochs=0,
    ################################
    device=None,
    callbacks=[
        ckpt_callback,
        trace_callback,
    ],
)

    Training

    You only need to simply specify the training phases in order when using the tools/train.py script, and the corresponding solver will be called automatically according to the training phase to execute the training process.

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

    • float: normal floating-point training.

    • qat: QAT training (Quantized Awareness Training), this stage first initializes a floating-point model, loads the trained floating-point model weights, and then converts this floating-point model into a QAT model for training.

    • int_infer: Fixed-point transformation prediction, this stage first initializes a floating-point model, converts the floating-point model into a QAT model and loads the trained QAT model weights, and then converts the QAT model into a fixed-point model. The converted fixed-point model cannot be trained and can only perform validation to obtain the final fixed-point model accuracy.

    Resume Training

    Unexpectedly interrupted training can be resumed by configuring the resume_optimizer and resume_epoch_or_step fields in {stage}_trainer of config, or by resuming only the optimizer for fine-tuning. For example:

    float_trainer = dict(
    ...
    model_convert_pipeline=dict(
        ...
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path="your_checkpoint_path",  # checkpoint path
            ),
        ],
    ),
    resume_optimizer=True,      # resume optimizer
    resume_epoch_or_step=True,  # resume epoch and step 
    ...,
)

    Training recovery has three scenarios:

    1. Full Recovery: This scenario is to resume the training that was unexpectedly interrupted, and will restore all the states of the previous checkpoint, including optimizer, LR, epoch, step, and so on. In this scenario, you only need to configure the resume_optimizer field.

    2. Resume Optimizer for Fine-tune: This scenario will only restore the state of optimizer and LR, with epoch and step reset to 0 for the fine-tuning of certain tasks. In this scenario, you need to configure both resume_optimizer and resume_epoch_or_step=False.

    3. Load Model Parameters Only: This scenario loads only model parameters and does not restore any other state (optimizeizer, epoch, step, or LR). In this scenario, you need to configure LoadCheckpoint in model_convert_pipeline, resume_optimizer=False, and resume_epoch_or_step=False.