yamle.quantization.qat package

class yamle.quantization.qat.QATQuantizer(learning_rate, epochs, *args, **kwargs)

Bases: BaseQuantizer

This is the quantization-aware training quantizer class.

It performs quantization-aware training on the model. In contrast to the static quantizer, the quantizer uses the calibration (validation) dataset to fine-tune the model. It uses the same optimiser as the one used for training the model.

Parameters:

• learning_rate (float) – The learning rate to use for the fine-tuning.

• epochs (int) – The number of epochs to use for the fine-tuning.

prepare(trainer, method)

This method is used to prepare the model for quantization.

It caches the original hyperparameters for the optimisation and replaces the hyperparameters with the ones for the fine-tuning.

Return type:

None

cleanup(method, trainer, *args, **kwargs)

This method is used to clean up the model after quantization.

Return type:

None

get_qconfig()

This method is used to get the quantization configuration.

We use the number of activation and weight bits to create the quantization configuration.

Return type:

Any

static add_specific_args(parent_parser)

This method is used to add specific arguments to the parser.

Return type:

ArgumentParser