neural_compressor.adaptor.mxnet_utils.util

Mxnet util module.

Classes

`OpType`	Enum op types.
`DataLoaderWrap`	DataLoader Wrap.
`DataIterLoader`	DataIterLoader.
`CollectorBase`	Collector Base class.
`CalibCollector`	Collect the calibration thresholds depending on the algorithm set.
`TensorCollector`	Tensors collector.
`NameCollector`	Name collector.
`CalibData`	Calibration data class.

Functions

`isiterable`(→ bool)	Checks whether object is iterable.
`ensure_list`(x)	Ensures that object is a list.
`check_mx_version`(version)	Checks MXNet version.
`combine_capabilities`(current, new)	Combine capabilities.
`make_nc_model`(target, sym_model, ctx, input_desc)	Converts a symbolic model to an Neural Compressor model.
`fuse`(sym_model, ctx)	Fuse the supplied model.
`get_framework_name`(ctx)	Get the framework name by version.
`prepare_model_data`(nc_model, ctx, data_x)	Prepares sym_model and dataloader needed for quantization, calibration or running.
`prepare_model`(nc_model, ctx, input_desc)	Prepare model.
`create_data_example`(ctx, input_desc)	Create data example by mxnet input description and ctx.
`prepare_dataloader`(nc_model, ctx, data_x)	Prepare dataloader.
`ndarray_to_device`(ndarray, device)	Ndarray to device.
`is_model_quantized`(sym_model)	Checks whether the model is quantized.
`query_quantizable_nodes`(sym_model, ctx, dataloader)	Query quantizable nodes of the given model.
`quantize_sym_model`(sym_model, ctx, qconfig)	Quantizes the symbolic model according to the configuration.
`run_forward`(sym_model, ctx, dataloader, b_filter[, ...])	Run forward propagation on the model.
`make_symbol_block`(sym_model, ctx, input_desc)	Convert a symbol model to gluon SymbolBlock.
`make_module`(sym_model, ctx, input_desc)	Convert a symbol model to Module.
`parse_tune_config`(tune_cfg, quantizable_nodes)	Convert the strategy config to MXNet quantization config.
`distribute_calib_tensors`(calib_tensors, calib_cfg, ...)	Distributes the tensors for calibration, depending on the algorithm set in the configuration of their nodes.
`calib_model`(qsym_model, calib_data, calib_cfg)	Calibrate the quantized symbol model using data gathered by the collector.
`amp_convert`(sym_model, input_desc, amp_cfg)	Convert model to support amp.

Module Contents

class neural_compressor.adaptor.mxnet_utils.util.OpType[source]: Enum op types.

neural_compressor.adaptor.mxnet_utils.util.isiterable(obj) → bool[source]

Checks whether object is iterable.

Parameters:: obj – object to check.
Returns:: True if object is iterable, else False.
Return type:: boolean

neural_compressor.adaptor.mxnet_utils.util.ensure_list(x)[source]

Ensures that object is a list.

Parameters:: x – input.
Returns:: x if x is list, else [x].
Return type:: list

neural_compressor.adaptor.mxnet_utils.util.check_mx_version(version)[source]

Checks MXNet version.

Parameters:: version (str) – version to check.
Returns:: True if mx.__version__ >= version, else False.
Return type:: boolean

neural_compressor.adaptor.mxnet_utils.util.combine_capabilities(current, new)[source]

Combine capabilities.

Parameters:

current (dict) – current capabilities.
new (dict) – new capabilities.

Returns:

contains all capabilities.

Return type:

dict

neural_compressor.adaptor.mxnet_utils.util.make_nc_model(target, sym_model, ctx, input_desc)[source]

Converts a symbolic model to an Neural Compressor model.

Parameters:

target (object) – target model type to return.
sym_model (tuple) – symbol model (symnet, args, auxs).
input_desc (list) – model input data description.

Returns:

converted neural_compressor model

Return type:

NCModel

neural_compressor.adaptor.mxnet_utils.util.fuse(sym_model, ctx)[source]

Fuse the supplied model.

Parameters:: sym_model (tuple) – symbol model (symnet, args, auxs).
Returns:: fused symbol model (symnet, args, auxs).
Return type:: tuple

neural_compressor.adaptor.mxnet_utils.util.get_framework_name(ctx)[source]

Get the framework name by version.

Parameters:: ctx (object) – mxnet context object.
Returns:: framework name.
Return type:: str

neural_compressor.adaptor.mxnet_utils.util.prepare_model_data(nc_model, ctx, data_x)[source]

Prepares sym_model and dataloader needed for quantization, calibration or running.

Parameters:

nc_model (object) – model to prepare.
data_x (object) – data iterator/loader to prepare.

Returns:

symbol model (symnet, args, auxs) and DataLoaderWrap.

Return type:

tuple

neural_compressor.adaptor.mxnet_utils.util.prepare_model(nc_model, ctx, input_desc)[source]

Prepare model.

Parameters:

nc_model (object) – model to prepare.
ctx (object) – mxnet context object.
input_desc (list) – input list of mxnet data types.

Returns:

mxnet model (symnet, args, auxs).

Return type:

object

neural_compressor.adaptor.mxnet_utils.util.create_data_example(ctx, input_desc)[source]

Create data example by mxnet input description and ctx.

Parameters:

ctx (object) – mxnet context object.
input_desc (list) – input list of mxnet data types.

Returns:

data example.

Return type:

list

neural_compressor.adaptor.mxnet_utils.util.prepare_dataloader(nc_model, ctx, data_x)[source]

Prepare dataloader.

Parameters:

nc_model (object) – model to prepare.
ctx (object) – mxnet context object.
data_x (object) – mxnet io iterable object or dataloader object.

Returns:

dataloader.

Return type:

object

neural_compressor.adaptor.mxnet_utils.util.ndarray_to_device(ndarray, device)[source]

Ndarray to device.

Parameters:

ndarray (ndarray) – model to prepare.
device (object) – mxnet device object.

Returns:

ndarray on the device.

Return type:

ndarray

neural_compressor.adaptor.mxnet_utils.util.is_model_quantized(sym_model)[source]

Checks whether the model is quantized.

Parameters:: sym_model (tuple) – symbol model (symnet, args, auxs).
Returns:: True if model is quantized, else False.
Return type:: boolean

neural_compressor.adaptor.mxnet_utils.util.query_quantizable_nodes(sym_model, ctx, dataloader)[source]

Query quantizable nodes of the given model.

Parameters:: sym_model (tuple) – symbol model (symnet, args, auxs) to query.
Returns:: quantizable nodes of the given model. dict: tensor to node mapping.
Return type:: list

neural_compressor.adaptor.mxnet_utils.util.quantize_sym_model(sym_model, ctx, qconfig)[source]

Quantizes the symbolic model according to the configuration.

Parameters:

sym_model (tuple) – symbol model (symnet, args, auxs).
qconfig (dict) – quantization configuration.

Returns:

Symbol model (symnet, args, auxs) and list of tensors for calibration.

Return type:

tuple

neural_compressor.adaptor.mxnet_utils.util.run_forward(sym_model, ctx, dataloader, b_filter, collector=None, pre_batch=None, post_batch=None)[source]

Run forward propagation on the model.

Parameters:

sym_model (tuple) – symbol model (symnet, args, auxs).
dataloader (DataLoaderWrap) – data loader.
b_filter (generator) – filter on which batches to run inference on.
collector (object) – collects information during inference.
pre_batch – function to call prior to batch inference.
post_batch – function to call after batch inference.

Returns:

batch count.

Return type:

int

neural_compressor.adaptor.mxnet_utils.util.make_symbol_block(sym_model, ctx, input_desc)[source]

Convert a symbol model to gluon SymbolBlock.

Parameters:

sym_model (tuple) – symbol model (symnet, args, auxs).
input_desc (list) – model input data description.

Returns:

SymbolBlock model.

Return type:

mx.gluon.SymbolBlock

neural_compressor.adaptor.mxnet_utils.util.make_module(sym_model, ctx, input_desc)[source]

Convert a symbol model to Module.

Parameters:

sym_model (tuple) – symbol model (symnet, args, auxs).
input_desc (list) – model input data description.

Returns:

Module model.

Return type:

mx.module.Module

neural_compressor.adaptor.mxnet_utils.util.parse_tune_config(tune_cfg, quantizable_nodes)[source]

Convert the strategy config to MXNet quantization config.

Parameters:

tune_cfg (dict) – tune config from neural_compressor strategy.
quantizable_nodes (list) – quantizable nodes in the model.

Returns:

quantization configuration. dict: calibration configuration.

Return type:

dict

neural_compressor.adaptor.mxnet_utils.util.distribute_calib_tensors(calib_tensors, calib_cfg, tensor_to_node)[source]

Distributes the tensors for calibration, depending on the algorithm set in the configuration of their nodes.

Parameters:

calib_tensors – tensors to distribute.
calib_cfg (dict) – calibration configuration.
tensor_to_node (dict) – tensor to node mapping.

Returns:

kl tensors and minmax tensors.

Return type:

tuple

neural_compressor.adaptor.mxnet_utils.util.calib_model(qsym_model, calib_data, calib_cfg)[source]

Calibrate the quantized symbol model using data gathered by the collector.

Parameters:

qsym_model (tuple) – quantized symbol model (symnet, args, auxs).
calib_data (CalibData) – data needed for calibration (thresholds).
calib_cfg (dict) – calibration configuration.

Returns:

quantized calibrated symbol model (symnet, args, auxs).

Return type:

tuple

neural_compressor.adaptor.mxnet_utils.util.amp_convert(sym_model, input_desc, amp_cfg)[source]: Convert model to support amp.

class neural_compressor.adaptor.mxnet_utils.util.DataLoaderWrap(dataloader, input_desc)[source]: DataLoader Wrap.

class neural_compressor.adaptor.mxnet_utils.util.DataIterLoader(data_iter)[source]: DataIterLoader.

class neural_compressor.adaptor.mxnet_utils.util.CollectorBase[source]: Collector Base class.

class neural_compressor.adaptor.mxnet_utils.util.CalibCollector(include_tensors_kl, include_tensors_minmax, num_bins=8001)[source]: Collect the calibration thresholds depending on the algorithm set.

class neural_compressor.adaptor.mxnet_utils.util.TensorCollector(include_nodes, qtensor_to_tensor, tensor_to_node)[source]

Tensors collector.

Builds up qtensor_to_tensor mapping.

class neural_compressor.adaptor.mxnet_utils.util.NameCollector[source]: Name collector.

class neural_compressor.adaptor.mxnet_utils.util.CalibData(cache_kl={}, cache_minmax={}, tensors_kl=[], tensors_minmax=[])[source]: Calibration data class.