AIMET PyTorch Bias Correction API

User Guide Link

To learn more about this technique, please see Cross-Layer Equalization

Bias Correction API

aimet_torch.bias_correction.correct_bias(model, quant_params, num_quant_samples, data_loader, num_bias_correct_samples, conv_bn_dict=None, perform_only_empirical_bias_corr=True, layers_to_ignore=None)[source]

Corrects bias for each Conv layer of model (unless ignored). A combination of Analytical and Empirical Bias Correction is used i.e. all the layers which can be corrected using Analytical Bias Correction are corrected using Analytical Bias Correction and remaining layers are corrected using Empirical method.

Returns an in-place corrected floating point model

Parameters:

model (Module) – Model to be corrected
quant_params (QuantParams) – Named tuple for quantization simulation for bias correction
num_quant_samples (int) – number of samples of images to pass through quantization sim for bias correction.
data_loader – data loader for the model
num_bias_correct_samples (int) – number of samples for Bias correction
conv_bn_dict (Optional[Dict[Module, ConvBnInfoType]]) – Dict of conv and bn with information related to activation. If None, the function calc it
perform_only_empirical_bias_corr (bool) – Default True. If true will perform only empirical Bias Corr for all layers irrespective of the fact that layer is eligible for Analytical Bias Corr.
layers_to_ignore (Optional[List[Module]]) – list of layer names for which we need to skip bias correction.

ConvBnInfoType

class aimet_common.bias_correction.ConvBnInfoType(input_bn=None, output_bn=None, in_activation_type=ActivationType.no_activation, out_activation_type=ActivationType.no_activation)[source]

Type for hoding convs with bn info and activation types Activation types supported are Relu and Relu6

Parameters:

input_bn – Reference to Input BatchNorm to layer
output_bn – Reference to Output BatchNorm to layer
in_activation_type (ActivationType) – Type of Activation
out_activation_type (ActivationType) – Type of Activation

ActivationType

class aimet_common.defs.ActivationType(value)[source]

Enums to identify activation type

no_activation = 0: No activation

relu = 1: ReLU activation

relu6 = 2: ReLU6 activation

Quantization Params

class aimet_torch.quantsim.QuantParams(weight_bw=8, act_bw=8, round_mode='nearest', quant_scheme=QuantScheme.post_training_tf_enhanced, config_file=None)[source]

Data type to hold quantization related params.

Constructor

Parameters:

weight_bw (int) – Weight bitwidth (4-31) to use for quantizing layer weights. Default = 8
act_bw (int) – Activation bitwidth(4-31) to use for quantizing layer activations. Default = 8
round_mode (str) – Rounding mode. Supported options are ‘nearest’ or ‘stochastic’
quant_scheme (Union[QuantScheme, str]) – Quantization scheme. Supported options are ‘tf_enhanced’ or ‘tf’ or using Quant Scheme Enum QuantScheme.post_training_tf or QuantScheme.post_training_tf_enhanced
config_file (Optional[str]) – Path to Configuration file for model quantizers

Code Example #1 Empirical Bias Correction

Load the model

    model = MobileNetV2()
    model.eval()

Apply Empirical Bias Correction

    from aimet_torch import bias_correction
    from aimet_torch.quantsim import QuantParams

    params = QuantParams(weight_bw=4, act_bw=4, round_mode="nearest", quant_scheme='tf_enhanced')

    # User action required
    # The following line of code is an example of how to use the ImageNet data's validation data loader.
    # Replace the following line with your own dataset's validation data loader.
    data_loader = ImageNetDataPipeline.get_val_dataloader()

    # Perform Empirical Bias Correction
    bias_correction.correct_bias(model.to(device="cuda"), params, num_quant_samples=1000,
                                 data_loader=data_loader, num_bias_correct_samples=512)

Code Example #2 Analytical + Empirical Bias correction

Load the model

    model = MobileNetV2()
    model.eval()

Find BN and Conv Modules

Find BN + Conv module pairs for analytical Bias Correction and remaining Conv modules for Empirical Bias Correction.

    module_prop_dict = bias_correction.find_all_conv_bn_with_activation(model, input_shape=(1, 3, 224, 224))

Apply Analytical + Empirical Bias Correction

    from aimet_torch import bias_correction
    from aimet_torch.quantsim import QuantParams

    params = QuantParams(weight_bw=4, act_bw=4, round_mode="nearest", quant_scheme='tf_enhanced')

    # User action required
    # The following line of code is an example of how to use the ImageNet data's validation data loader.
    # Replace the following line with your own dataset's validation data loader.
    data_loader = ImageNetDataPipeline.get_val_dataloader()

    # Perform Bias Correction
    bias_correction.correct_bias(model.to(device="cuda"), params, num_quant_samples=1000,
                                 data_loader=data_loader, num_bias_correct_samples=512,
                                 conv_bn_dict=module_prop_dict, perform_only_empirical_bias_corr=False)