neural_compressor.torch.utils.utility

Module Contents

Classes

Mode

Generic enumeration.

Functions

register_algo(name)

Decorator function to register algorithms in the algos_mapping dictionary.

fetch_module(model, op_name)

Get module with a given op name.

set_module(model, op_name, new_module)

Set module with a given op name.

get_quantizer(model, quantizer_cls[, quant_config])

Get the quantizer.

postprocess_model(model, mode, quantizer)

Process quantizer attribute of model according to current phase.
neural_compressor.torch.utils.utility.register_algo(name)[source]

Decorator function to register algorithms in the algos_mapping dictionary.

Usage example:

@register_algo(name=example_algo) def example_algo(model: torch.nn.Module, quant_config: RTNConfig) -> torch.nn.Module:

…

Parameters:

name (str) – The name under which the algorithm function will be registered.

Returns:

The decorator function to be used with algorithm functions.

Return type:

decorator

neural_compressor.torch.utils.utility.fetch_module(model, op_name)[source]

Get module with a given op name.

Parameters:

  • model (object) – the input model.

  • op_name (str) – name of op.

Returns:

module (object).

neural_compressor.torch.utils.utility.set_module(model, op_name, new_module)[source]

Set module with a given op name.

Parameters:

  • model (object) – the input model.

  • op_name (str) – name of op.

  • new_module (object) – the input model.

Returns:

module (object).

class neural_compressor.torch.utils.utility.Mode[source]

Generic enumeration.

Derive from this class to define new enumerations.

neural_compressor.torch.utils.utility.get_quantizer(model, quantizer_cls, quant_config=None, *args, **kwargs)[source]

Get the quantizer.

Initialize a quantizer or get quantizer attribute from model.

Parameters:

  • model (torch.nn.Module) – pytorch model.

  • quantizer_cls (Quantizer) – quantizer class of a specific algorithm.

  • quant_config (dict, optional) – Specifies how to apply the algorithm on the given model. Defaults to None.

Returns:

quantizer object.

neural_compressor.torch.utils.utility.postprocess_model(model, mode, quantizer)[source]

Process quantizer attribute of model according to current phase.

In prepare phase, the quantizer is set as an attribute of the model to avoid redundant initialization during convert phase.

In ‘convert’ or ‘quantize’ phase, the unused quantizer attribute is removed.

Parameters:

  • model (torch.nn.Module) – pytorch model.

  • mode (Mode) – The mode of current phase, including ‘prepare’, ‘convert’ and ‘quantize’.

  • quantizer (Quantizer) – quantizer object.