create_quant_retrain_config

Applicability

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

Description

Finds all quantizable layers in a graph, creates a quantization configuration file, and writes the quantization configuration of the quantizable layers to the configuration file.

Prototype

create_quant_retrain_config(config_file, model, input_data, config_defination=None)

Parameters

Parameter	Input/Output	Description
config_file	Input	Path (including the file name) of the QAT configuration file. The existing file (if any) in the path will be overwritten upon this API call. A string.
model	Input	Original model for QAT, with weights loaded. Data type: torch.nn.Module
input_data	Input	Input data of the model. A torch.tensor is replaced with an equivalent tuple(torch.tensor). A tuple.
config_defination	Input	Simplified configuration file. The simplified configuration file quant.cfg is generated based on the retrain_config_pytorch.proto file. The .proto file is stored in /amct_pytorch/proto/* under the AMCT installation directory. For details about the parameters in the .proto file and the generated simplified quantization configuration file quant.cfg, see Simplified QAT Configuration File. Default: None* A string.

Returns

None

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)])
 
# Create a quantization configuration file.
amct.create_quant_retrain_config(config_file="./configs/config.json",
                            model=model,
                            input_data=input_data)

Flush file: an INT8 QAT configuration file in JSON format. The following is an example. (The configuration file output by this API will be overwritten when QAT is performed again.) For details about the parameters, see Quantization Configuration File.

{
    "version":1,
    "batch_num":1,
    "conv1":{
        "retrain_enable":true,
        "retrain_data_config":{
            "algo":"ulq_quantize",
            "dst_type":"INT8"
        },
        "retrain_weight_config":{
            "algo":"arq_retrain",
            "channel_wise":true,
            "dst_type":"INT8"
        }
    },
    "layer1.0.conv1":{
        "retrain_enable":true,
        "retrain_data_config":{
            "algo":"ulq_quantize",
            "dst_type":"INT8"
        },
        "retrain_weight_config":{
            "algo":"arq_retrain",
            "channel_wise":true,
            "dst_type":"INT8"
        }
    },
    "fc":{
        "retrain_enable":true,
        "retrain_data_config":{
            "algo":"ulq_quantize",
            "dst_type":"INT8"
        },
        "retrain_weight_config":{
            "algo":"arq_retrain",
            "channel_wise":false,
            "dst_type":"INT8"
        }
    }
}

Parent topic: QAT APIs