save_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

Applies to static compression combination. Based on the retrained model, generates a fake-quantized model for accuracy simulation and a deployable model, which have undergone static compression combination.

Prototype

1
save_compressed_retrain_model(model, record_file, save_path, input_data, input_names=None, output_names=None, dynamic_axes=None)

Parameters

Parameter

Input/Output

Description

model

Input

PyTorch model after static compression combination.

A torch.nn.Module.

record_file

Input

Path (including the file name) of the sparsity and quantization factor record file.

A string.

save_path

Input

Path for storing the compressed 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

Names of the model input, which are displayed in the saved quantized ONNX model.

Default: None

A list of strings.

output_names

Input

Names of the model output, which are displayed in the saved quantized 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

Restrictions

If the model is sparsified only (not quantized), 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.

Example

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import amct_pytorch as amct
# Build a graph of the network for compression.
model = build_model()

# create compressed model

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

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

# Insert the API for saving the compressed model and convert it into an ONNX file.
amct.save_compressed_retrain_model(
     compressed_retrain_model,
     record_file,
     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 files, including *deploy.external and *fakequant.external:

    This type of file is generated only when the size of the saved fake-quantized model and deployable model file is greater than or equal to 2 GB. The *.external file is generated in the same directory as the compressed *.onnx model file and is used to save the data in the tensor. Each tensor data is saved in a separate .external file. The file name is the same as the tensor name, for example, conv1.weight_deploy.external and conv1.weight_fakequant.external.

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

When static compression combination is performed again, the preceding files output by the API will be overwritten.