auto_channel_prune_search

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

Calculates the sparsity sensitivity (affecting accuracy) and sparsity gain (affecting performance) of each channel based on the user model. Then, the search policy searches for the optimal layer-wise channel sparsity ratio based on the input to balance accuracy and performance. Finally, a configuration file is generated.

Prototype

auto_channel_prune_search(model, config, input_data, output_cfg, sensitivity, search_alg)

Parameters

Parameter	Input/Output	Description
model	Input	PyTorch model to be sparsified. A torch.nn.Module.
config	Input	Path of the auto channel pruning configuration file. The simplified configuration file is generated based on AutoChannelPruneConfig in the basic_info.proto file. The .proto file is stored in /amct_pytorch/proto/* under the AMCT installation directory. For details about the parameters in the *.proto file and the generated auto channel pruning search configuration file, see Simplified Configuration File for Auto Channel Pruning Search. A string.
input_data	Input	Input data (including labels) provided by the user. A list[data,label]. The data type of the list element is torch.tensor.
output_cfg	Input	Path of the output channel pruning configuration file. A string.
sensitivity	Input	Sensitivity calculation method. A string or subclass of SensitivityBase. The string is an existing method of AMCT. Currently, the value can be TaylorLossSensitivity. The subclass of SensitivityBase can be instantiated and defined by the user.
search_alg	Input	Method of searching for channels to be sparsified. A string or subclass of SearchChannelBase. The string is an existing method of AMCT. Currently, the value can be GreedySearch. The subclass of SearchChannelBase can be instantiated and defined by the user.

Returns

None

Example

import amct_pytorch as amct
# Construct the input data.
input_data = torch.randn(input_shape)         
model.eval()        
output = model.forward(input_data)        
labels = torch.randn(output.size())        
data = [input_data,labels]

amct.auto_channel_prune_search(
     model=model,
     config='./tmp/sample.cfg',
     input_data=data,  
     output_cfg='./tmp/output.cfg', 
     sensitivity='TaylorLossSensitivity', 
     search_alg='GreedySearch')

Flush file:

Auto channel pruning configuration file. This file needs to be transferred to the channel sparsity API for subsequent services.

Parent topic: Auto Channel Pruning Search API