create_prune_retrain_model
Applicability
Product |
Supported |
|---|---|
Atlas 350 Accelerator Card |
|
|
|
|
|
|
|
|
|
|
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 configuration file, inserts or replaces related operators in the input graph, generates a sparsity record file record_file, and returns a resultant torch.nn.Module model that can be used for retraining.
Prototype
1 | prune_retrain_model = create_prune_retrain_model (model, input_data, config_defination, record_file) |
Parameters
Parameter |
Input/Output |
Description |
|---|---|---|
model |
Input |
Model to be sparsified, 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 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. |
record_file |
Input |
Path (including the file name) of the file that records sparsity information. This file records the cascade correlations between filter-level sparsity nodes or 2:4 structured sparsity nodes. A string. |
Returns
Returns the resultant torch.nn.Module model that can be used for post-sparsity training.
Example
1 2 3 4 5 6 7 8 9 10 11 12 13 14 | import amct_pytorch as amct # Build a graph of the network for pruning. model = build_model() model.load_state_dict(torch.load(state_dict_path)) input_data = tuple([torch.randn(input_shape)]) # Call the API for sparsifying models. record_file = os.path.join(TMP, 'scale_offset_record.txt') cfg_file = './prune_config.cfg' prune_retrain_model = amct.create_prune_retrain_model( model, input_data, cfg_file, record_file) |