Data Calibration

In quantization-aware training (QAT), one critical step is determining the quantization parameter scale. A reasonable scale can significantly improve the training results of the model and accelerate its convergence. Calibration involves running a few batches of data from the training set through the floating-point model (only the forward process, without backward propagation), collecting the distribution histogram of the data, and calculating the min_value and max_value using specific methods. These min_value and max_value can then be used to obtain the scale. When the training accuracy of QAT cannot be improved, performing calibration to fine-tune the quantization parameters and obtain the scale before starting QAT can provide better initialization parameters for quantization, improving both convergence speed and accuracy.

How to Define a Calibration Model

  • By default, no modifications to the existing model structure are required.

    Calibration uses a Float2Calibration Converter to convert the model from a floating-point model to a quantized model.

An example is:

convert_mode = "fuse_bn"
float2calibration = dict(
                type="Float2Calibration",
                convert_mode=convert_mode,
                example_data=example_data,
                qconfig_setter=calibration_qconfig_setter,
            )

The example_data must have the same format as the real input data of the model, and users can obtain example_data from the data_loader.
Float2Calibration also supports directly passing a dataloader: example_data_loader=calibration_data_loader.
For the configuration of qconfig_setter, users can refer to the section Qconfig Configuration.

Performing Calibration with a Floating-Point Model

HAT integrates the Calibration functionality. The command for performing Calibration with a floating-point model is similar to normal training. Simply execute the following command:

python3 tools/train.py --stage calibration ...

It is important to note some configurations in the calibration_trainer section of the config file:

# Note: The transforms of the dataset during calibration can be
# consistent with that during training or validation, or customized.
# Default used `val_batch_processor`.
calibration_data_loader = copy.deepcopy(data_loader)
calibration_batch_processor = copy.deepcopy(val_batch_processor)

calibration_trainer = dict(
    type="Calibrator",
    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="Float2Calibration"),
        ],
    ),
    # 2. Set data_loader 和 batch_processor
    data_loader=calibration_data_loader,
    batch_processor=calibration_batch_processor,
    num_stages=30,
    ...   
)

1. Dataset Setup:

The dataset used for Calibration cannot be the test set (it can be the training set or other data). During Calibration, the data augmentation transforms can be consistent with those used during normal training, or they can be set to match the transforms used during validation. Custom transforms can also be defined. (There is no definitive conclusion on which approach yields the best experimental results, so all options can be tried.)

2. Number of Images for Calibration Iterations (For Reference):

  • Classification: Generally, using 500 to 1500 images can achieve good results.
  • Segmentation & Detection: Using around 100 to 300 images is sufficient.
Note

The specific number of images is not fixed. The above suggestions are based on experience from existing experiments and can be adjusted according to actual needs.

Using the Calibration Model for QAT Training

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, "calibration-checkpoint-best.pth.tar"
                ),
            ),
        ],
    ),
)

Setting the averaging_constant Parameter for QAT:

The update rule for the quantization scale parameter is:
scale = (1 - averaging_constant) * scale + averaging_constant * current_scale.

In some existing experiments (mainly image classification tasks), it was found that after performing calibration, fixing the activation scale without updating it (i.e., setting the activation averaging_constant=0) and setting the weight averaging_constant=1 may yield slightly better results.

Note

This configuration is not suitable for all tasks. For example, in LiDAR tasks, fixing the scale may lead to a decrease in accuracy. Adjustments should be made based on the actual situation.

Next, simply execute the normal QAT training command to start QAT training:

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