#Quantization Training Process

This document only explains the operations required for performing quantization training in HAT. For the basic principles of quantization and its implementation in the training framework, please refer to the relevant documentation of horizon_plugin_pytorch.

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

• Model Registration: All modules in HAT adopt a registration mechanism. Only models registered in the corresponding registry can be used in the config file in the form of dict(type={$class_name}, ...).

• qconfig Configuration: Before quantizing the model, we need to configure the quantization settings (qconfig) for the model. horizon_plugin_pytorch provides the QconfigSetter interface to quickly configure the model's qconfig. For details, refer to Qconfig Configuration.

Additionally, to enable the model to be converted into a quantized model, certain conditions must be met. For details, refer to the relevant documentation of horizon_plugin_pytorch.

#Adding a Custom Model

import torch
from torch import nn

from hat.registry import OBJECT_REGISTRY


# Registering the Model Using a Decorator
@OBJECT_REGISTRY.register_module
class ExampleNet(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        ...

#Adding a 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=...,
    ...,
)

#Training

Simply specify the training stages in order when using the tools/train.py script. The corresponding solver will be automatically called to execute the training process based on the specified stage:

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

• float: Standard floating-point training.

• qat: Quantization-Aware Training (QAT). First, initialize a floating-point model, load the trained weights of the floating-point model, and then convert the model to a QAT model for training. The output of the training is a pseudo-quantized QAT model.

#Resuming Training

You can resume interrupted training or fine-tune by configuring the resume_optimizer and resume_epoch_or_step fields in the {stage}_trainer section of the config. For example:

float_trainer = dict(
    ...
    model_convert_pipeline=dict(
        ...
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path="your_checkpoint_path",  # Path to the checkpoint to resume
            ),
        ],
    ),
    resume_optimizer=True,      # Resume optimizer
    resume_epoch_or_step=True,  # Resume epoch or step
    ...,
)

There are three scenarios for resuming training:

1. Full Recovery: This scenario is used to resume training after an unexpected interruption. It restores all states from the last checkpoint, including the optimizer, learning rate (LR), epoch, step, etc. In this case, you only need to configure the resume_optimizer field.

2. Resume Optimizer for Fine-Tuning: This scenario only restores the states of the optimizer and LR, while the epoch and step will start from 0. It is used for fine-tuning specific tasks. In this case, you need to configure resume_optimizer and set resume_epoch_or_step=False.

3. Load Model Parameters Only: This scenario only loads the model parameters without restoring any other states (optimizer, epoch, step, LR). In this case, you only need to configure LoadCheckpoint in the model_convert_pipeline and set resume_optimizer=False and resume_epoch_or_step=False.