save_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

Generates the final ONNX fake-quantized model and deployable model based on the retrained sparse model.

Restrictions

For sparse models, the two files generated by this API are the ONNX files exported using PyTorch. The file content is the same, and the file names contain the deploy and fake_quant keywords, respectively.

Prototype

save_prune_retrain_model(model, save_path, input_data, input_names=None, output_names=None, dynamic_axes=None)

Parameters

Parameter	Input/Output	Description
model	Input	PyTorch model that has been sparsified. A torch.nn.Module.
save_path	Input	Path for storing the compressed model. Must include the prefix of the model name, for example, **./prune_model/model***. A string.
input_data	Input	Input data of the model. A torch.tensor is replaced with an equivalent tuple(torch.tensor). A tuple.
input_names	Input	Model input name, which is displayed in the saved sparse ONNX model. Default: None A list of strings.
output_names	Input	Model output name, which is displayed in the saved sparse ONNX model. Default: None A list of strings.
dynamic_axes	Input	Dynamic axes of the model inputs and outputs. For example, if the inputs have format NCHW, where N, H and W are uncertain, and the outputs have format NL, where N is uncertain, then pass: {"inputs": [0,2,3], "outputs": [0]}, where 0, 2, and 3 indicate the indexes of N, H, and W, respectively. Default: None A dict<string, dict<python:int, string>>, or dict<string, list(int)>.

Returns

None

Example

import amct_pytorch as amct
# Build a graph of the network for compression.
model = build_model()

# create selective prune model

# Train the retrained model to calculate quantization factors.
train(pruned_retrain_model)
infer(pruned_retrain_model)

input_data = tuple([torch.randn(input_shape)])
save_path = os.path.join(OUTPUTS_DIR, 'custom_name')

# Insert the API for saving the compressed model and convert it into an ONNX file.
amct.save_prune_retrain_model(
     pruned_retrain_model,
     save_path,
     input_data,
     input_names=['input'],
     output_names=['output'],
     dynamic_axes={'input':{0:'batch_size'}, 'output':{0:'batch_size'}})

Flush files:

A fake-quantized ONNX model file for accuracy simulation on ONNX Runtime with the file name containing the fake_quant keyword.
A deployable ONNX model file with the file name containing the deploy keyword. The model can be deployed on the AI processor after being converted by ATC.
(Optional) External data file. The sparsity feature does not distinguish between the *deploy.external and *fakequant.external files.
Only the size of the saved fake-quantized model and deployable model file is available. This type of file is generated only when the file size is greater than or equal to 2 GB and is generated in the same directory as the compressed *.onnx model file. It is used to store the data in the tensor. Each tensor data is stored in a separate file with the same file name as the tensor, for example, conv_1.weight.

When ATC is used to load the compressed *.onnx deployable model file for model conversion, the tensor data in the external data file in the same directory is automatically read.

When sparsity is performed again, the preceding files output by the API will be overwritten.

Parent topic: Sparsity APIs