save_distill_model

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	√
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	√
Atlas inference products	√
Atlas training products	√

Description

Generates a fake-quantized model for accuracy simulation and a deployable model based on the distilled model.

Prototype

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

Parameters

Parameter	Input/Output	Description
model	Input	Quantized model after distillation. A torch.nn.Module.
save_path	Input	Path for saving the distilled model. Must include the prefix of the model name, for example, **./quantized_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.
record_file	Input	Path (including the file name) of the quantization factor record file. Default: None A string. Restrictions: If None, the quantization factor record file is stored in the amct_log folder.
input_names	Input	Names of the model input nodes, which are displayed in the saved quantized ONNX model. Default: None A list of strings.
output_names	Input	Names of the model output nodes, 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 an input has format NCHW, where N, H and W are dynamic, and an output has format NL, where N is dynamic, then: {"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 distillation.
model = build_model()
model.load_state_dict(torch.load(state_dict_path))
input_data = tuple([torch.randn(input_shape)])
 
# Insert the distillation API and save the distilled model as an ONNX file.
amct.save_distill_model(
               model, 
               "./model/distilled",
               input_data,
               record_file="./results/records.txt",
               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.

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

Parent topic: Distillation APIs