#Quantized Awareness Training

After completing Calibration, if the current quantization results still cannot meet the accuracy requirements, the next step is to perform Quantization Aware Training (QAT).

Compared with pure floating-point model training, QAT usually makes the training process more challenging because of the additional quantization constraints introduced during training. The goal of the QAT toolkit is to reduce training difficulty while also lowering the engineering effort required for deploying quantized models.

The objective of this stage is to further optimize quantization error based on the existing quantization parameter initialization, thereby providing a higher-accuracy QAT model for subsequent model export, fixed-point conversion, and deployment validation.

#BN Strategy and Training Mode

Before starting QAT, it is recommended to determine the BN handling strategy and the current training mode first. These settings directly affect the training behavior and also affect the accuracy and convergence patterns you will observe later.

#FuseBN

FuseBN means that after Prepare and before QAT training, the BN weight and bias are absorbed into the Conv weight and bias, and BN is no longer updated independently during training. This absorption process itself is lossless, and it is also the default QAT training method.

Suitable when:

• The model itself is suitable for BN fusion.

• You want the training path and the final deployment path to stay as aligned as possible.

#WithBN

WithBN means that Conv + BN stays unfused during QAT training, BN is trained together with the rest of the model, and fusion is performed only when the deployment model is generated after training.

Suitable when:

• The floating-point training stage already uses FreezeBN or a similar strategy.

• You want the QAT stage to align more closely with the floating-point training method.

#Related Interface

qat_mode is used to control the current quantization training mode. If you need to switch BN strategy or training mode, it is recommended to decide based on the current model structure, floating-point training strategy, and the later deployment target, instead of applying a training habit directly.

The related interfaces are defined in horizon_plugin_pytorch.qat_mode. According to qat_mode.py, three modes are currently supported, and the default value is QATMode.FuseBN.

The most commonly used interfaces here are:

• QATMode: used to specify the concrete quantization training mode. The available values are FuseBN, WithBN, and WithBNReverseFold.

• set_qat_mode: used to set the current quantization training mode, and it is usually recommended to call it before prepare.

• get_qat_mode: used to check which quantization training mode is currently active.

In practice, the usual choice is:

• If there is no special requirement, start from the default FuseBN.

• If float training already uses FreezeBN or a similar strategy and you want to preserve the BN behavior during QAT, try WithBN first.

• If you clearly plan to absorb BN gradually during QAT and want a smaller disturbance before and after absorption, then consider WithBNReverseFold.

For example, if the floating-point training stage uses FreezeBN or a similar strategy and you want BN to remain involved during QAT, you can configure it as follows:

from horizon_plugin_pytorch.qat_mode import QATMode, get_qat_mode, set_qat_mode
from horizon_plugin_pytorch.quantization import prepare

# The default is usually FuseBN
print("current qat_mode:", get_qat_mode())

# Set the QAT mode before prepare
set_qat_mode(QATMode.WithBN)

# If you plan to absorb BN gradually later in QAT,
# you can also switch to QATMode.WithBNReverseFold
# set_qat_mode(QATMode.WithBNReverseFold)

qat_model = prepare(
    float_model,
    example_input,
    qconfig_setter=my_qconfig_setter,
)

If you still need to decide which qat_mode to use, or want to continue tuning it together with learning rate, scale update strategy, and high-precision configuration, continue with the Precision Tuning Guide, especially Common Mixed Precision Tuning Workflow, Basic Tuning Methods, and Advanced Tuning Methods.

#Quantization Parameter Update Strategy

Before training starts, whether the quantization parameters are fixed or continuously updated will also directly affect the later convergence behavior.

When Calibration accuracy is already good, fixing the feature map quantization parameters for QAT training can usually produce better results. When Calibration accuracy is poor, it is usually not recommended to directly fix the quantization parameters obtained from Calibration. In that case, it is better to go back and first resolve the quantization bottleneck in the Calibration stage.

If the quantization parameters from Calibration still do not meet the accuracy target and you want them to keep updating during QAT, the tool provides the following update methods.

#Statistics-based Method

For performance reasons, it is recommended to use MinMaxObserver during the QAT stage to continuously collect statistics and update the quantization parameters. For details, refer to QConfig Details.

#Learning-based Method (Not Recommended)

The tool also supports the LSQ algorithm (see Learned Step Size Quantization). The corresponding interface is _LearnableFakeQuantize. For details, refer to Definition of FakeQuantize.

#Workflow and Example

Although the Quantized Awareness Training tool does not strictly require you to provide a pretrained floating-point model at the beginning, in practice it is still recommended to start from a pretrained high-accuracy floating-point model, because this usually makes the training process much easier.

from horizon_plugin_pytorch.quantization import FakeQuantState, prepare, set_fake_quantize
from horizon_plugin_pytorch.quantization import hbdk4 as hb4
from horizon_plugin_pytorch.quantization.qconfig_setter import *
from hbdk4.compiler import convert

# Convert the model to QAT state
qat_model = prepare(
    float_model,
    example_input,
    qconfig_setter=QconfigSetter(
        get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
        ],
    ),
).to(device)

# Load the quantization parameters from the Calibration model
qat_model.load_state_dict(calib_model.state_dict())

# Run quantization aware training
# As a finetuning process, QAT usually uses a smaller learning rate
optimizer = torch.optim.SGD(
    qat_model.parameters(), lr=0.0001, weight_decay=2e-4
)

for nepoch in range(epoch_num):
    # Note the way the QAT model training state is controlled here
    qat_model.train()
    set_fake_quantize(qat_model, FakeQuantState.QAT)

    train_one_epoch(
        qat_model,
        nn.CrossEntropyLoss(),
        optimizer,
        None,
        train_data_loader,
        device,
    )

    # Note the way the QAT model eval state is controlled here
    qat_model.eval()
    set_fake_quantize(qat_model, FakeQuantState.VALIDATION)

    # Evaluate the QAT model
    top1, top5 = evaluate(
        qat_model,
        eval_data_loader,
        device,
    )
    print(
        "QAT model: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
            top1.avg, top5.avg
        )
    )

    # Evaluate the quantized model
    qat_hbir_model = hb4.export(
        qat_model, example_input
    )
    quantized_hbir_model = convert(qat_hbir_model)

    top1, top5 = evaluate(
        quantized_hbir_model,
        eval_data_loader,
    )
    print(
        "Quantized model: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
            top1.avg, top5.avg
        )
    )

The training loop above is only the basic form. In real usage, QAT is usually not finished in a single round. Instead, you typically need to use staged training results and validation results to judge whether the current configuration is appropriate, and whether you still need to continue training or adjust the strategy.

Therefore, besides the training itself, two things should be done continuously:

1. Judge from the current results whether the learning rate, BN strategy, quantization parameter update method, or high-precision configuration still needs adjustment.

2. Do not only look at the QAT model accuracy. Continue validating the quantized model accuracy as well, so that you can confirm whether the current training result is truly beneficial for later deployment.

From the example above, compared with conventional floating-point training, QAT mainly adds two key steps:

1. prepare: transform the floating-point network and insert fake quantization nodes.

2. Load the parameters from the Calibration model: by loading the pseudo-quantization parameters obtained during Calibration, the model gets a better initialization.

For state_dict processing, you should not only pay attention to keys and values; _metadata also needs to be preserved. The following is an example of copying a state_dict:

new_state_dict = OrderedDict()
for k, v in state_dict.items():
    new_state_dict[k] = v

if hasattr(state_dict, "_metadata"):
    new_state_dict._metadata = copy.deepcopy(state_dict._metadata)

At this point, the fake-quantized model construction and parameter initialization are finished. You can then continue with normal training iterations and parameter updates while continuously monitoring the quantized model accuracy.

To support segmented deployment or align with the floating-point training strategy, you may also need to freeze some parts of the model during training. The following is a common way to do that:

from horizon_plugin_pytorch.quantization import freeze_qat_module

# Model weights / quantization parameters will be frozen, and all operators will be set to eval state.
# Make sure freeze_qat_module is called after interfaces such as train() / eval() / set_fake_quantize() that may change model state.
freeze_qat_module(model)

#What to Monitor During Training

During QAT training, it is recommended to monitor at least the following metrics together.

#QAT model accuracy

QAT model accuracy is mainly used to observe the overall convergence trend during training. It helps you judge whether the training is still moving in a better direction, but it should not be used alone to judge the final deployment accuracy.

#quantized model accuracy

quantized model accuracy is closer to the actual result of the later deployment path, and is used to judge whether the current training result is truly helpful for later export, fixed-point conversion, and deployment.

If you only look at the QAT model accuracy and ignore the quantized model accuracy, it becomes easier to run into the situation where “training looks fine, but deployment accuracy drops” in later export or board verification stages.

#Whether the loss is close to the final float training loss

QAT is usually a finetuning process on top of an already converged floating-point model, so whether the loss stays close to the final float training loss can help you judge whether quantization training has entered a normal convergence state.

• If the loss is too large, it often means the model is hard to converge.

• If the loss is too small, it may also introduce generalization issues.

#Whether NAN / INF / long-term non-convergence appears

These phenomena usually indicate issues in the training pipeline itself, scale updates, quantization configuration, or initialization. They should be investigated as early as possible instead of waiting until the end of training.

#Common Issues and Recommended Actions