aihwkit.nn.low_precision_conversion module

Functions to convert a given model to a quantized counterpart

aihwkit.nn.low_precision_conversion.convert_to_quantized(model, quantization_map)[source]

High level function to perform the quantization of a model according to the QuantizationMap defined by the user. See the QuantizationMap dataclass for instructions on how to define its fields.

This function calls the recursive quantize_model function which performs the actual conversion between the individual modules. It then enables the quantization states for all the quantized modules.

Parameters:

  • model (Module) – The model to quantize

  • quantization_map (QuantizationMap) – The dataclass that contains instructions on how to quantize the model

Returns:

The quantized model

Return type:

torch.nn.Module

aihwkit.nn.low_precision_conversion.quantize_model(model, quantization_map, model_name='')[source]

Traverses a model recursively and replaces a module with a quantized counterpart, if such a conversion is defined in the quantization map. It realizes all the capabilities of the QuantizationMap, namely: - Excluding specific modules from the quantization, identified by exact state_dict

string (quantization_map.excluded_modules)

  • Instance specific quantization, identified by exact state_dict string

    (quantization_map.instance_qconfig_map)

  • Module specific quantization, identified by the type of a module

    (quantization_map.module_qconfig_map)

  • Input quantization on specified modules, identified by the state_dict string

    (quantization_map.input_activation_qconfig_map)

(see QuantizationMap and the related examples for more details on how to define these fields)

Parameters:

  • model (Module) – Model to recursively quantize

  • quantization_map (QuantizationMap) – The dataclass that defines the quantization conversions as well as the quantization parameters for each conversion

  • model_name (str, optional) – The name of the current module in the state dict of the original model, by default “”

Returns:

Quantized model

Return type:

Module