create_distill_model

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	√
Atlas inference product	√
Atlas training product	√

Description

Quantizes the input graph structure based on the specified distillation configuration file, inserts quantization-related operators (distillation layer and searching N layer of activations and weights) into the input graph structure, and returns the modified torch.nn.Module model that can be used for distillation.

Prototype

compress_model = create_distill_model(config_file, model, input_data)

Parameters

Parameter	Input/Output	Description
config_file	Input	User-generated distillation configuration file, which is used to specify the configuration and distillation structure of the quantization layer in the model network. A string. Restrictions: The config.json file passed to this API must be the same as that passed to the create_distill_config API.
model	Input	Floating-point original model to be distilled, 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.

Parameter

Input/Output

Description

config_file

Input

User-generated distillation configuration file, which is used to specify the configuration and distillation structure of the quantization layer in the model network.

A string.

Restrictions: The config.json file passed to this API must be the same as that passed to the create_distill_config API.

model

Input

Floating-point original model to be distilled, 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.

Returns

Returns the resultant torch.nn.Module model that can be used for distillation.

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)])

# Generate a compressed model.
compress_model = amct.create_distill_model(
                 config_json_file,
                 model,
                 input_data)

Parent topic: Distillation APIs