quantize_preprocess

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	√
Atlas inference product	√
Atlas training product	√

Description

Preprocesses the quantization of a graph based on the quantization configuration file, inserts the balanced quantization operators, generates a balanced quantization factor record file record_file, and returns a torch.nn.Module model ready for calibration.

Prototype

calibration_model = quantize_preprocess(config_file, record_file, model, input_data)

Parameters

Parameter	Input/Output	Description
config_file	Input	User-generated distillation configuration file, which is used to specify the configuration of the quantization layer in the model network. A string.
record_file	Input	Path (including the file name) of the balanced quantization factor record file. A string.
model	Input	Original model for quantization, with weights loaded. A torch.nn.Module.
input_data	Input	Input data of the model. A torch.tensor is replaced with an equivalent tuple(torch.tensor). A tuple.

Returns

Returns the resultant torch.nn.Module model for calibration.

Example

import amct_pytorch as amct
# Build a graph of the network to be quantized.
model = build_model()
model.load_state_dict(torch.load(state_dict_path))
input_data = tuple([torch.randn(input_shape)])

tensor_balance_factor_record_file = os.path.join(TMP, 'tensor_balance_factor_record.txt')
modified_model = os.path.join(TMP, 'modified_model.onnx')
# Insert the quantization API.
calibration_model = amct.quantize_preprocess(config_json_file,
                                             tensor_balance_factor_record_file,
                                             model,
                                             input_data)

Parent topic: PTQ APIs