distill

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

Distills the input graph structure based on the given distillation configuration file and returns the modified torch.nn.Module model.

Prototype

distill_model = distill(model, compress_model, config_file, train_loader, epochs=1, lr=1e-3, sample_instance=None, loss=None, optimizer=None)

Parameters

Parameter	Input/Output	Description
model	Input	Floating-point original model to be distilled, with weights loaded. A torch.nn.Module.
compress_model	Input	Resultant torch.nn.Module model that can be used for distillation. A torch.nn.Module. Restrictions: The model passed to this API must be a compressed model.
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.
train_loader	Input	Training dataset. torch.utils.data.DataLoader Restrictions: It must match the model input size.
epochs	Input	Maximum number of epochs. Default: 1 An int.
lr	Input	Learning rate. Default: 1e-3 A float.
sample_instance	Input	Instantiated object of the method provided by the user for obtaining model input data. Default: None DistillSampleBase Restrictions: It must be inherited from the DistillSampleBase class and implement the get_model_input_data method. For details, see the /amct_pytorch/distill/distill_sample.py file under the AMCT installation directory.
loss	Input	Instantiated object for computing the loss. Default: None torch.nn.Modules.loss._Loss
optimizer	Input	Instantiated object of the optimizer. Default: None torch.optim.Optimizer

Returns

Returns the resultant torch.nn.Module model 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))
compress_model = compress(model)
input_data = tuple([torch.randn(input_shape)])
train_loader = torch.utils.data.DataLoader(input_data)
loss = torch.nn.MSELoss()
optimizer = torch.optim.AdamW(compress_model.parameters(), lr=0.1)

# Perform distillation.
distill_model = amct.distill(
                model,
                compress_model
                config_json_file,
                train_loader,
                epochs=1,
                lr=1e-3,
                sample_instance=None, 
                loss=loss,
                optimizer=optimizer)

Parent topic: Distillation APIs