4.2.4.1. Quantization API

horizon_plugin_pytorch.quantization.fuse_modules.fuse_modules(model, modules_to_fuse, inplace=False, fuser_func=None, fuse_custom_config_dict=None)

Fuse a list of modules into a single module.

Fuses only the following sequence of modules: conv, bn; conv, bn, relu; conv, relu; conv, bn, add; conv, bn, add, relu; conv, add; conv, add, relu; linear, bn; linear, bn, relu; linear, relu; linear, bn, add; linear, bn, add, relu; linear, add; linear, add, relu. For these sequences, the first element in the output module list performs the fused operation. The rest of the elements are set to nn.Identity().

Parameters

  • model – Model containing the modules to be fused.

  • modules_to_fuse – List of list of module names to fuse. Can also be a list of strings if there is only a single list of modules to fuse.

  • inplace – Bool specifying if fusion happens in place on the model, by default a new model is returned.

  • fuser_func – Function that takes in a list of modules and outputs a list of fused modules of the same length. For example, fuser_func([convModule, BNModule]) returns the list [ConvBNModule, nn.Identity()]. Defaults to torch.ao.quantization.fuse_known_modules.

  • fuse_custom_config_dict –

    Custom configuration for fusion:

    fuse_custom_config_dict = {
    "additional_fuser_method_mapping": {
        (torch.nn.Conv2d, torch.nn.BatchNorm2d): fuse_conv_bn
    },
}

Returns

Model with fused modules. A new copy is created if inplace=False.

Return type

module

Examples

>>> m = M().eval()
>>> # m is a module containing the sub-modules below
>>> modules_to_fuse = [ ['conv1', 'bn1', 'relu1'],
                      ['submodule.conv', 'submodule.relu']]
>>> fused_m = fuse_modules(
                m, modules_to_fuse)
>>> output = fused_m(input)

>>> m = M().eval()
>>> # Alternately provide a single list of modules to fuse
>>> modules_to_fuse = ['conv1', 'bn1', 'relu1']
>>> fused_m = fuse_modules(
                m, modules_to_fuse)
>>> output = fused_m(input)
horizon_plugin_pytorch.quantization.quantize.convert(module, mapping=None, inplace=False, remove_qconfig=True, fast_mode=False)

Convert modules.

Convert submodules in input module to a different module according to mapping by calling from_float method on the target module class. And remove qconfig at the end if remove_qconfig is set to True.

Parameters

  • module – Input module.

  • mapping – A dictionary that maps from source module type to target module type, can be overwritten to allow swapping user defined Modules.

  • inplace – Carry out model transformations in-place, the original module is mutated.

  • remove_qconfig – Remove qconfig at the end if True.

  • fast_mode – Whether to accelerate quantized model forward. If set True, quantized model cannot be compiled.

horizon_plugin_pytorch.quantization.quantize.prepare_calibration(model, mapping=None, inplace=False)

Prepare a model for calibration (eager mode).

Parameters

  • model – Input float model.

  • mapping (dict) – Optional float-to-quantized module mapping.

  • inplace (bool) – Carry out model transformations in-place.

horizon_plugin_pytorch.quantization.quantize.prepare_qat(model, mapping=None, inplace=False, optimize_graph=False, hybrid=False, optimize_kwargs=None)

Prepare qat.

Prepare a copy of the model for quantization-aware training and converts it to quantized version.

Quantization configuration should be assigned preemptively to individual submodules in .qconfig attribute.

Parameters

  • model – Input model to be modified in-place.

  • mapping (dict) – Dictionary that maps float modules to quantized modules to be replaced.

  • inplace (bool) – Carry out model transformations in-place, the original module is mutated.

  • optimize_graph (bool) – Whether to do some process on origin model for special purpose. Currently only support using torch.fx to fix cat input scale (only used on Bernoulli).

  • hybrid (bool) –

    Whether to generate a hybrid model that some intermediate operation is computed in float. There are some constraints for this functionality now:

    1. The hybrid model cannot pass check_model and cannot be compiled.

    2. Some quantized operation cannot directly accept input from float operation, user need to manually insert QuantStub.

  • optimize_kwargs (dict) –

    A dict for optimize graph, for example:

    optimize_kwargs = {
    "opt_types": ("unify_inputs_scale",),
    "module_prefixes": ("backbone.conv",),
    "module_types": (horizon.nn.qat.conv2d,),
    "functions": (torch.clamp,),
    "methods": ("add",),
}

horizon_plugin_pytorch.quantization.quantize_fx.convert_fx(graph_module, inplace: bool = False, convert_custom_config_dict: Optional[Dict[str, Any]] = None, _remove_qconfig: bool = True, fast_mode: bool = False)

Convert a calibrated or trained model to a quantized model.

Parameters

  • graph_module – A prepared and calibrated/trained model (GraphModule).

  • inplace (bool) – Carry out model transformations in-place, the original module is mutated.

  • convert_custom_config_dict (dict) –

    Dictionary for custom configurations for convert function:

    convert_custom_config_dict = {
    # We automatically preserve all attributes, this option is
    # just in case and not likely to be used.
    "preserved_attributes": ["preserved_attr"],
}

  • _remove_qconfig (bool) – Option to remove the qconfig attributes in the model after convert. For internal use only.

  • fast_mode (bool) – Whether to accelerate quantized model forward. If set True, quantized model cannot be compiled.

Returns

A quantized model (GraphModule).

Return type

QuantizedGraphModule

Examples

# prepared_model: the model after prepare_fx/prepare_qat_fx and
# calibration/training
quantized_model = convert_fx(prepared_model)
horizon_plugin_pytorch.quantization.quantize_fx.fuse_fx(model, fuse_custom_config_dict: Optional[Dict[str, Any]] = None)

Fuse modules like conv+add+bn+relu etc.

Fusion rules are defined in horizon_plugin_pytorch.quantization.fx.fusion_pattern.

Parameters

  • model – A torch.nn.Module model.

  • fuse_custom_config_dict (dict) –

    Dictionary for custom configurations for fuse_fx, e.g.:

    fuse_custom_config_dict = {
    # We automatically preserve all attributes, this option is
    # just in case and not likely to be used.
    "preserved_attributes": ["preserved_attr"],
}

Returns

Fused model.

Return type

GraphModuleWithAttr

Examples

from torch.quantization import fuse_fx
m = fuse_fx(m)
horizon_plugin_pytorch.quantization.quantize_fx.prepare_calibration_fx(model, qconfig_dict=None, prepare_custom_config_dict=None)

Prepare a model for calibration (fx mode).

Parameters

  • model – Input float model or GraphModule.

  • qconfig_dict (dict) – Qconfig rules, same schema as prepare_qat_fx.

  • prepare_custom_config_dict (dict) – Customization configuration dictionary.

horizon_plugin_pytorch.quantization.quantize_fx.prepare_qat_fx(model, qconfig_dict=None, prepare_custom_config_dict=None, optimize_graph=False, hybrid=False, hybrid_dict=None)

Prepare a model for quantization aware training.

Parameters

  • model – torch.nn.Module model or GraphModule model (maybe from fuse_fx).

  • qconfig_dict (dict) –

    Dictionary with the following configurations:

    qconfig_dict = {
    "": qconfig,
    "module_type": [
        (torch.nn.Conv2d, qconfig),
    ],
    "module_name": [
        ("foo.bar", qconfig),
    ],
}
# priority (in increasing order):
#   global, module_type, module_name, module.qconfig
# qconfig == None means quantization should be skipped.

  • prepare_custom_config_dict (dict) –

    Customization configuration dictionary for quantization tool:

    prepare_custom_config_dict = {
    "preserved_attributes": ["preserved_attr"],
}

  • optimize_graph (bool) – Whether to do some process on origin model for special purpose. Currently only support using torch.fx to fix cat input scale (only used on Bernoulli).

  • hybrid (bool) –

    Whether prepare model in hybrid mode. Default is False and model runs on BPU completely. It should be True if the model is quantized by model convert or contains some CPU ops. In hybrid mode, ops which aren’t supported by BPU and ops which are specified by the user will run on CPU.

    How to set qconfig: Qconfig in hybrid mode is the same as qconfig in non-hybrid mode. For BPU op, we should ensure the input of this op is quantized, the activation qconfig of its previous non-quantstub op should not be None even if its previous non-quantstub op is a CPU op.

    How to specify CPU op: Define CPU module_name or module_type in hybrid_dict.

  • hybrid_dict (dict) –

    Dictionary to define user-specified CPU op:

    hybrid_dict = {
    "module_type": [torch.nn.Conv2d, ...],
    "module_name": ["foo.bar", ...],
}
# priority (in increasing order): module_type, module_name
# To set a function or method as CPU op, wrap it as a module.

Returns

A GraphModule with fake quant modules (configured by qconfig_dict), ready for quantization aware training.

Return type

ObservedGraphModule

Examples

import torch
from horizon_plugin_pytorch.quantization import get_default_qat_qconfig
from horizon_plugin_pytorch.quantization import prepare_qat_fx

qconfig = get_default_qat_qconfig()
def train_loop(model, train_data):
    model.train()
    for image, target in data_loader:
        ...

qconfig_dict = {"": qconfig}
prepared_model = prepare_qat_fx(float_model, qconfig_dict)
train_loop(prepared_model, train_loop)
horizon_plugin_pytorch.fx.fx_helper.wrap(skip_compile: bool = False)

Extend torch.fx.wrap.

This function can be:

  1. called or used as a decorator on a string to register a builtin function as a “leaf function”

  2. called or used as a decorator on a function to register this function as a “leaf function”

  3. called or used as a decorator on subclass of torch.nn.Module to register this module as a “leaf module”, and register all user defined method in this class as “leaf method”

  4. called or used as a decorator on a class method to register it as “leaf method”

Parameters

skip_compile (bool) – Whether the wrapped part should not be compiled.

Returns

The actual decorator.

Return type

wrap_inner

class horizon_plugin_pytorch.quantization.fake_quantize.FakeQuantize(observer=None, saturate=None, in_place=False, compat_mask=True, channel_len=1, **observer_kwargs)

Simulate the quantize and dequantize operations in training time.

The output of this module is given by:

x_out = (clamp(round(x/scale + zero_point), quant_min, quant_max)
         - zero_point) * scale

scale defines the scale factor used for quantization.

zero_point specifies the quantized value to which 0 in floating point maps to.

quant_min specifies the minimum allowable quantized value.

quant_max specifies the maximum allowable quantized value.

fake_quant_enabled controls the application of fake quantization on tensors, note that statistics can still be updated.

observer_enabled controls statistics collection on tensors.

dtype specifies the quantized dtype that is being emulated with fake-quantization, the allowable values is qint8 and qint16. The values of quant_min and quant_max should be chosen to be consistent with the dtype.

Parameters

  • observer (type) – Module for observing statistics on input tensors and calculating scale and zero-point.

  • saturate (bool) – Whether zero out the grad for value out of quanti range.

  • in_place (bool) – Whether use in place fake quantize.

  • compat_mask (bool) – Whether pack the bool mask into bitfield when saturate = True.

  • channel_len (int) – Size of data at channel dim.

  • observer_kwargs – Arguments for the observer module.

extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

set_qparams(scale, zero_point=None)

Set qparams, default symmetric.

classmethod with_args(**kwargs)

Wrapper that allows creation of class factories.

This can be useful when there is a need to create classes with the same constructor arguments, but different instances.

Examples

>>> Foo.with_args = classmethod(_with_args)
>>> foo_builder = Foo.with_args(a=3, b=4).with_args(answer=42)
>>> foo_instance1 = foo_builder()
>>> foo_instance2 = foo_builder()
>>> id(foo_instance1) == id(foo_instance2)
False
horizon_plugin_pytorch.quantization.fake_quantize.default_8bit_fake_quant(*args, **kwargs)

Default 8-bit fake quantize for activations.

horizon_plugin_pytorch.quantization.fake_quantize.default_weight_8bit_fake_quant(*args, **kwargs)

Default 8-bit fake quantize for weights.

class horizon_plugin_pytorch.quantization.observer.MovingAverageMinMaxObserver(averaging_constant=0.01, dtype=None, qscheme=None, quant_min=None, quant_max=None, is_sync_quantize=False, factory_kwargs=None)

MovingAverageMinMax Observer.

Observer module for computing the quantization parameters based on the moving average of the min and max values.

This observer computes the quantization parameters based on the moving averages of minimums and maximums of the incoming tensors. The module records the average minimum and maximum of incoming tensors, and uses this statistic to compute the quantization parameters.

Parameters

  • averaging_constant – Averaging constant for min/max.

  • dtype – Quantized data type.

  • qscheme – Quantization scheme to be used, only support per_tensor_symmetric scheme.

  • reduce_range – Reduces the range of the quantized data type by 1 bit.

  • quant_min – Minimum quantization value.

  • quant_max – Maximum quantization value.

  • is_sync_quantize – Whether use sync quantize.

  • factory_kwargs – Arguments for register data buffer.

forward(x_orig)

Record the running minimum and maximum of x.

class horizon_plugin_pytorch.quantization.observer.MovingAveragePerChannelMinMaxObserver(averaging_constant=0.01, ch_axis=0, dtype=None, qscheme=None, quant_min=None, quant_max=None, is_sync_quantize=False, factory_kwargs=None)

MovingAveragePerChannelMinMax Observer.

Observer module for computing the quantization parameters based on the running per channel min and max values.

This observer uses the tensor min/max statistics to compute the per channel quantization parameters. The module records the running minimum and maximum of incoming tensors, and uses this statistic to compute the quantization parameters.

Parameters

  • averaging_constant – Averaging constant for min/max.

  • ch_axis – Channel axis.

  • dtype – Quantized data type.

  • qscheme – Quantization scheme to be used, only support per_channel_symmetric.

  • quant_min – Minimum quantization value.

  • quant_max – Maximum quantization value.

  • is_sync_quantize – Whether use sync quantize.

  • factory_kwargs – Arguments for register data buffer.

forward(x_orig)

Defines the computation performed at every call.

horizon_plugin_pytorch.quantization.fuse_modules.fuse_known_modules(mod_list, is_qat=False, additional_fuser_method_mapping=None)

Fuse modules.

Return a list of modules that fuses the operations specified in the input module list.

Fuses only the following sequence of modules: conv, bn; conv, bn, relu; conv, relu; conv, bn, add; conv, bn, add, relu; conv, add; conv, add, relu; linear, bn; linear, bn, relu; linear, relu; linear, bn, add; linear, bn, add, relu; linear, add; linear, add, relu. For these sequences, the first element in the output module list performs the fused operation. The rest of the elements are set to nn.Identity().

class horizon_plugin_pytorch.march.March

BPU platform.

BAYES: Bayes platform

BERNOULLI2: Bernoulli2 platform

horizon_plugin_pytorch.quantization.qconfig.get_default_calib_qconfig(activation_fake_quant='fake_quant', weight_fake_quant='fake_quant', activation_observer='percentile', weight_observer='min_max', activation_qkwargs=None, weight_qkwargs=None)

Get default calibration qconfig.

Same parameters as get_default_qconfig, with activation_observer="percentile" by default.

horizon_plugin_pytorch.quantization.qconfig.get_default_qat_out_qconfig(activation_fake_quant=None, weight_fake_quant='fake_quant', activation_observer=None, weight_observer='min_max', activation_qkwargs=None, weight_qkwargs=None)

Get default QAT output qconfig.

Same parameters as get_default_qconfig, but activation fake quant and observer default to None so the last layer can keep a more accurate float/high-precision output. Typical usage is model.classifier.qconfig = get_default_qat_out_qconfig().

horizon_plugin_pytorch.quantization.qconfig.get_default_qat_qconfig(activation_fake_quant='fake_quant', weight_fake_quant='fake_quant', activation_observer='min_max', weight_observer='min_max', activation_qkwargs=None, weight_qkwargs=None)

Get default QAT qconfig.

Same parameters as get_default_qconfig. Typical usage is model.qconfig = get_default_qat_qconfig() before prepare_qat or prepare_qat_fx.

horizon_plugin_pytorch.quantization.check_model(module, example_inputs, march=None, input_source='ddr', advice=None)

Check if nn.Module or jit.ScriptModule can be compiled by HBDK.

Dump advices for improving performance on BPU.

Parameters

  • module – nn.Module or jit.ScriptModule.

  • example_inputs (tuple) – Example inputs in torch.tensor format, for jit.trace and shape inference.

  • march (str) – Target BPU march. Valid options are bayes and bernoulli2. If not provided, use horizon plugin global march.

  • input_source (str or sequence) – Input feature sources (ddr / resizer / pyramid).

  • advice (int) – Print HBDK compiler advices for improving the utilization of the model on BPU if layers of the model become slow by more than the specified time (in microseconds).

Returns

0 if pass, otherwise not.

Return type

int

horizon_plugin_pytorch.quantization.compile_model(module, example_inputs, hbm, march=None, name=None, input_source='ddr', input_layout=None, output_layout='NCHW', opt='O2', balance_factor=2, progressbar=True, jobs=16, debug=True, extra_args=None)

Compile the nn.Module or jit.ScriptModule.

Parameters

  • module – nn.Module or jit.ScriptModule.

  • example_inputs (tuple) – Example inputs in torch.tensor format.

  • hbm (str) – Output path of hbdk-cc.

  • march (str) – Target BPU march (bayes / bernoulli2).

  • name (str) – Name of the model, recorded in hbm.

  • input_source (str or sequence) – Input feature sources (ddr / resizer / pyramid).

  • input_layout (str) – NHWC, NCHW or BPU_RAW.

  • output_layout (str) – NHWC, NCHW or BPU_RAW.

  • opt (str or int) – O0, O1, O2, O3, ddr, fast, balance.

  • balance_factor (int) – Balance ratio when opt is ‘balance’.

  • progressbar (bool) – Show compilation progress.

  • jobs (int) – Compiler threads. 0 means all available hardware concurrency.

  • debug (bool) – Enable debugging info in hbm.

  • extra_args (list) – Extra args listed in hbdk-cc -h.

Returns

0 if pass, otherwise not.

Return type

int

horizon_plugin_pytorch.quantization.export_hbir(module, example_inputs, hbir, march=None)

Export the nn.Module or jit.ScriptModule to hbdk3.HBIR.

Parameters

  • module – nn.Module or jit.ScriptModule.

  • example_inputs (tuple) – Example inputs in torch.tensor format.

  • hbir (str) – Output path of hbir.

  • march (str) – Target march (bayes / bernoulli2).

Returns

Input names and output names.

Return type

tuple

horizon_plugin_pytorch.quantization.perf_model(module, example_inputs, march=None, out_dir='.', name=None, hbm=None, input_source='ddr', input_layout=None, output_layout='NCHW', opt='O3', balance_factor=2, progressbar=True, jobs=16, layer_details=False, extra_args=None)

Estimate the performance of nn.Module or jit.ScriptModule.

Parameters

  • module – nn.Module or jit.ScriptModule.

  • example_inputs (tuple) – Example inputs in torch.tensor format.

  • march (str) – Target BPU march (bayes / bernoulli2).

  • out_dir (str) – Output directory for performance results.

  • name (str) – Name of the model, recorded in hbm.

  • hbm (str) – Output path of hbdk-cc.

  • input_source (str or sequence) – Input feature sources (ddr / resizer / pyramid).

  • input_layout (str) – NHWC, NCHW or BPU_RAW.

  • output_layout (str) – NHWC, NCHW or BPU_RAW.

  • opt (str or int) – O0, O1, O2, O3, ddr, fast, balance.

  • balance_factor (int) – Balance ratio when opt is ‘balance’.

  • progressbar (bool) – Show compilation progress.

  • jobs (int) – Compiler threads.

  • layer_details (bool) – Show layer performance details (dev use only).

  • extra_args (list) – Extra args listed in hbdk-cc -h.

Returns

Performance details in json dict, or error code when fail.

Return type

dict or int

horizon_plugin_pytorch.quantization.visualize_model(module, example_inputs, march=None, save_path=None, show=True)

Visualize nn.Module or jit.ScriptModule at the view of HBDK.

Parameters

  • module – nn.Module or jit.ScriptModule.

  • example_inputs (tuple) – Example inputs in torch.tensor format, for jit.trace and shape inference.

  • march (str) – Target BPU march. Valid options are bayes and bernoulli2. If not provided, use horizon plugin global march.

  • save_path (str) – Path to save the plot image.

  • show (bool) – Display the plotted image via display. Make sure X-server is correctly configured.