restore_compressed_retrain_model

Applicability

Product	Supported
Atlas 350 Accelerator Card	QAT INT8 quantization: √ Filter-level sparsity: √ 2:4 structured sparsity: x
Atlas A3 training product/Atlas A3 inference product	QAT INT8 quantization: √ Filter-level sparsity: √ 2:4 structured sparsity: √
Atlas A2 training product/Atlas A2 inference product	QAT INT8 quantization: √ Filter-level sparsity: √ 2:4 structured sparsity: √
Atlas 200I/500 A2 inference product	QAT INT8 quantization: √ Filter-level sparsity: √ 2:4 structured sparsity: √
Atlas inference product	QAT INT8 quantization: √ Filter-level sparsity: √ 2:4 structured sparsity: x
Atlas training product	QAT INT8 quantization: √ Filter-level sparsity: √ 2:4 structured sparsity: x

Note: For the products marked with x, no error is reported when the API is called, but no performance gains are obtained.

Description

Compresses the input model to be statically combined based on the specified compression combination configuration file and record file (sparsification before quantization), and loads the saved weight. Sparsifies the input model based on the sparsity records in the given record_file, and then inserts quantization-related operators (QAT layer for activations and weights and searchN layer) into the model. Loads the checkpoint weight parameters saved during training and returns the modified torch.nn.Module model.

Prototype

compressed_retrain_model = restore_compressed_retrain_model(model, input_data, config_defination, record_file, pth_file, state_dict_name=None)

Parameters

Parameter	Input/Output	Description
model	Input	Torch model. 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.
config_defination	Input	Simplified configuration file for static compression combination. The simplified configuration file compressed.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 compressed*.cfg, see Simplified QAT Configuration File. A string.
record_file	Input	File that records the sparsity and quantization factors A string.
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 static compression combination training model torch.nn.Module that is sparsified based on the sparsity relationship in record_file, with quantization layers inserted and the weight file loaded.

Restrictions

The compression combination configuration file must contain at least one of the following configurations: sparsity configuration or quantization configuration.

Example

import amct_pytorch as amct
# Build a graph of the network for compression combination.
model = build_model()
input_data = tuple(torch.randn(input_shape))
save_pth_path = /your/path/to/save/tmp.pth

record_file = os.path.join(TMP, 'compressed_record.txt')
config_defination = './compressed_cfg.cfg'

torch.save({'state_dict': model.state_dict()}, save_pth_path)
compressed_retrain_model = amct.restore_compressed_retrain_model(
                                model,
                                input_data,
                                config_defination,
                                record_file,
                                save_pth_path,
                                'state_dict')

Parent topic: Compression Combination APIs