restore_prune_retrain_model

Applicability

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

Note: For the products marked with x, no error is reported when the API is called for the 2:4 structured sparsity feature, but no performance gains are obtained.

Description

Filter-level sparsity or 2:4 structured sparsity API. Only either of the two sparsity features can be enabled at a time. This API sparsifies the input graph based on the given sparsity record file record_file, and returns a resultant torch.nn.Module model that can be used for retraining.

Prototype

prune_retrain_model = restore_prune_retrain_model (model, input_data, record_file, config_defination, pth_file, state_dict_name=None)

Parameters

Parameter	Input/Output	Description
model	Input	Model to be sparsified, 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.
record_file	Input	Path (including the file name) of the sparsity record file, which is generated by the create_prune_retrain_model API, to ensure that the models generated by the two APIs are consistent. A string.
config_defination	Input	Simplified configuration file. The simplified configuration file prune.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 prune*.cfg, see Simplified QAT Configuration File. 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 resultant torch.nn.Module model that can be used for post-sparsity training.

Example

import amct_pytorch as amct
# Build a graph of the network for pruning.
config_defination = './prune_cfg.cfg'
model = build_model()
input_data = tuple([torch.randn(input_shape)])
save_pth_path = /your/path/to/save/tmp.pth
model.load_state_dict(torch.load(state_dict_path))
 
# Call the API for sparsifying models.
record_file = os.path.join(TMP, 'scale_offset_record.txt')
prune_retrain_model = amct.restore_prune_retrain_model(           
               model,
               input_data,
               record_file,
               config_defination,
               save_pth_path,
              'state_dict')

Parent topic: Sparsity APIs