#QAT Mode

In the section of Quantization Training Process - Adding the Config File mentioned above, we have already introduced how to configure the qat_mode. Here, we will provide a detailed introduction to the function and usage method of qat_mode.

Qat_mode is used to set whether to perform the quantization training with BN in the QAT phase. With the help of the FuseBN interface provided by HAT, it can also control whether to perform the training with BN throughout the whole process or with BN being gradually absorbed midway.

#Optional Definitions

The following three settings are available for qat_mode：

class QATMode(object):
    FuseBN = "fuse_bn"
    WithBN = "with_bn"
    WithBNReverseFold = "with_bn_reverse_fold"

#Principles

#Fuse BN

QAT Phase without BN (default quantization training method of HAT)

By setting qat_mode to fuse_bn, in the op fusion process of the floating-point model, the weight and bias of BN are absorbed into that of Conv, and the original combination of Conv + BN will be left with only Conv, and this absorption process is theoretically error-free.

#With BN

QAT Phase with BN

By setting qat_mode to with_bn, when the floating-point model is converted to QAT model, BN is not absorbed into Conv, but exists in the quantized model as a fused quantized op in the QAT phase in the form of Conv + BN + Output Quantized Node. Finally, at the end of quantization training, in the step where the model is converted to quantized (also called int infer), the weight and bias of BN will be automatically absorbed into the quantization parameters of Conv, where the quantized op obtained after the absorption remains consistent with the original QAT op calculation result.

In this mode, you can also choose to absorb the BN into Conv in the middle of QAT. The reason why the forward results of the QAT model before and after manually absorbing the BN are inconsistent is that after the BN weight is absorbed into the Conv weight, the quantized parameter conv_weight_scale calculated in the previous quantization training is no longer applicable to the current conv_weight and will lead to large errors in the quantization of conv_weight, which requires more quantization training and more adjustments on quantization parameters.

#With BN Reverse Fold

QAT Phase with BN

The difference between this mode and with_bn is that, in this mode, the BN weight is considered when calculating conv_weight_scale in the quantization training phase before the BN is absorbed (calculations are not detailed here), so that after absorbing the BN weight, the conv_weight_scale is still applicable to the new conv_weight.

This mode is intended to provide a lossless way of absorbing BNs step by step: absorbing BNs in the middle of the quantization training, the forward result of the model is theoretically identical before and after the absorption, and you can gradually absorb all the BNs in the model before the end of quantization training and ensure that the loss will not fluctuate too much after each absorption.

In this mode, if there are BNs not absorbed at the end of the quantization training, they will be automatically absorbed when the model is converted from QAT to quantized. In theory, such absorption is lossless.

#Usage

#Set qat_mode

Only needs to set qat_mode in model_convert_pipeline.

For example:

model_convert_pipeline=dict(
    type="ModelConvertPipeline",
    qat_mode="with_bn",
    converters=[
        dict(type="Float2QAT"),
        dict(
            type="LoadCheckpoint",
            checkpoint_path=os.path.join(
                ckpt_dir, "qat-checkpoint-best.pth.tar"
            ),
        ),
    ],
)

#View Current qat_mode

from horizon_plugin_pytorch.qat_mode import get_qat_mode
qat_mod = get_qat_mode()

#Set Progressive Absorption BN

In both with_bn and with_bn_reverse_fold modes, you can set FuseBN as a callback function to absorb the BN in the specified module at the specified epoch or step.

FuseBN definition:

class FuseBN(OnlineModelTrick):
    Args:
        module: sub model names to fuse BN.
        step_or_epoch: when to fuseBN, same length as module.
        update_by: by step or by epoch.
        inplace: if fuse BN inplace
    def __init__(
        self,
        modules: List[List[str]],
        step_or_epoch: List[int],
        update_by: str,
        inplace: bool = False,
    )

Use the FuseBN example in the config file:

from hat.callbacks import FuseBN

# Define the callback function
# BN in the backbone module is absorbed at the 1000th step
# BN in the neck module is absorbed at the 1500th step
fuse_bn_callback = FuseBN(
   modules=[['backbone'], ['neck']],
   step_or_epoch=[1000, 1500],
   update_by='step',
)

# Add the callback function to the trainer
qat_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=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"
                ),
            ),
        ],
    ),
    data_loader=...,
    optimizer=...,
    batch_processor=...,
    num_epochs=...,
    device=None,
    callbacks=[
        callbacks0,
        ..., 
        fuse_bn_callback,
        callbacks99
    ],
    ...,
)

#Qat_mode Summary

In general, a training process starts from the floating-point training, and when the desired accuracy is met, move on to the quantization training, where only fuse_bn is used. Only when the floating-point training is skipped, i.e., it starts with the quantization training, the quantization training mode with BN is needed to ensure the model converges.