#Quantized Awareness Training

The quantized awareness training is performed by inserting some pseudo-quantized nodes into the model, so as to minimize the loss of accuracy when the model obtained through quantized awareness training is converted into a fixed-point model. The quantized awareness training is no different from traditional model training in that one can start from scratch, build a pseudo-quantized model, and then train on that pseudo-quantized model. Due to the limitations of the deployed hardware platform, it is challenging to understand these limitations and build a pseudo-quantization model based on them. The quantized awareness training tool reduces the challenges of developing quantized models by automatically inserting pseudo-quantization operators into the provided floating-point model based on the limitations of the deployment platform.

The quantized awareness training is generally more difficult than the training of pure floating-point models due to the various restrictions imposed. The goal of the quantized awareness training tool is to reduce the difficulty of quantized awareness training and to reduce the engineering difficulty of quantized model deployment.

#Process and Example

Although our quantized awareness training tool does not mandate that you provide a pre-trained floating-point model at the outset, experience has shown that starting quantized awareness training from a pre-trained high-precision floating-point model generally significantly reduces the difficulty of training.

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 in the Calibration model
qat_model.load_state_dict(calib_model.state_dict())
# perform quantized awareness training
# as a filetune process, quantized awareness training generally requires setting a small learning rate
optimizer = torch.optim.SGD(
    qat_model.parameters(), lr=0.0001, weight_decay=2e-4
)

for nepoch in range(epoch_num):
    # note the method of controlling the training state of the QAT model 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 method of controlling the eval state of the QAT model here
    qat_model.eval()
    set_fake_quantize(qat_model, FakeQuantState.VALIDATION)

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

# test quantized model accuracy
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
    )
)

As can be seen from the above sample code, there are two additional steps in quantized awareness training compared to traditional pure floating-point model training:

1. prepare. The goal of this step is to transform the floating-point network and insert pseudo-quantized nodes.

2. Load the Calibration model parameters. A better initialization is obtained by loading the pseudo-quantization parameters obtained from Calibration.

   Modifications of state_dict should not only focus on the keys and values, but also on the _metadata. Below 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 construction of the pseudo-quantized model and the initialization of the parameters are completed, and then the regular training iterations and model parameter updates can be performed, and the quantized model accuracy can be monitored.

To meet the requirements of segmented deployment or to align with float training strategies, it may be necessary to freeze certain parts of the model during training. You can refer to the following code to perform the freezing:

from horizon_plugin_pytorch.quantization import freeze_qat_module

# Model weights / quantization parameters will be fixed, and all operators will be set to eval mode.
# Ensure that the freeze_qat_module interface is called after invoking interfaces like train(), eval(), or set_fake_quantize(), which may change the model state.
freeze_qat_module(model)

#Quantization Parameter Update Strategy

When the calibration accuracy is satisfactory, fixing the quantization parameters of the feature map for QAT training can yield better results. When the accuracy is poor, the quantization parameters obtained from calibration usually cannot be fixed (in this case, it is recommended to prioritize debugging the calibration accuracy to resolve quantization bottlenecks). If it is confirmed that the quantization parameters from the calibration phase cannot meet the accuracy requirements and you wish to continuously update them during QAT training, the tool provides the following update methods.

#Statistics-based Method

For performance reasons, we recommend using MinMaxObserver during the QAT phase to continuously track statistics and update quantization parameters. For usage instructions, please refer to Build QConfig.

#Learning-based Method (Not Recommended)

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