restore_quant_retrain_model

Applicability

Product	Supported
Atlas A3 training series products/Atlas A3 inference series products	INT8 quantization: √
Atlas A2 training products/Atlas A2 inference products	INT8 quantization: √
Atlas 200I/500 A2 inference product	INT8 quantization: √
Atlas inference series products	INT8 quantization: √
Atlas training products	INT8 quantization: √

Description

Quantizes a graph based on the given configuration file, inserts quantization-related layers (quantization-aware layers of activations and weights and layers for searching for N), generates a quantization factor record file (record_file), loads the checkpoint weight parameters saved during training, and returns the resultant model of the torch.nn.Module type.

Prototype

quant_retrain_model = restore_quant_retrain_model (config_file, model, record_file, input_data, pth_file, state_dict_name=None)

Parameters

Parameter	Input/Output	Description
config_file	Input	User-generated QAT configuration file, which is used to specify the configuration of the quantization layer in the model network. A string. Restrictions: The config.json file passed to this API must be the same as that passed to the create_quant_retrain_model API.
model	Input	Original model, with weights unloaded. A torch.nn.Module.
record_file	Input	Path (including the file name) of the quantization factor record file. A string.
input_data	Input	Input data of the model. A torch.tensor is replaced with an equivalent tuple(torch.tensor). A tuple.
pth_file	Input	Weight file saved during training. A string.
state_dict_name	Input	Key value corresponding to the weight in the weight file. Default: None A string.

Returns

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

Example

import amct_pytorch as amct
# Build a graph of the network to be quantized.
model = build_model()
input_data = tuple([torch.randn(input_shape)])
 
scale_offset_record_file = os.path.join(TMP, 'scale_offset_record.txt')
# Insert the quantization API.
quant_retrain_model = amct.restore_quant_retrain_model(
               config_json_file,
               model,
               scale_offset_record_file,
               input_data,
               pth_file)

Parent topic: QAT APIs