auto_decomposition

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

(Optional) Performs tensor decomposition on the input PyTorch model object to obtain the model object after decomposition and the names of the layers before and after decomposition, and saves the decomposition information file.

Prototype

model, changes = auto_decomposition(model, decompose_info_path=None)

Parameters

Parameter	Input/Output	Restriction
model	Input	PyTorch model object that contains pre-trained weights to be decomposed. When calling this API, you are advised to place the model on the CPU instead of the GPU to prevent insufficient GPU memory during decomposition. A torch.nn.Module.
decompose_info_path	Input	Path for storing the decomposition information file. The file is stored in JSON format. Therefore, the JSON file name extension is recommended. If the value is None, the decomposition information file is not saved (default). A string. Default: None

Parameter

Input/Output

Restriction

model

Input

PyTorch model object that contains pre-trained weights to be decomposed. When calling this API, you are advised to place the model on the CPU instead of the GPU to prevent insufficient GPU memory during decomposition.

A torch.nn.Module.

decompose_info_path

Input

Path for storing the decomposition information file. The file is stored in JSON format. Therefore, the JSON file name extension is recommended. If the value is None, the decomposition information file is not saved (default).

A string.

Default: None

Returns

Returns the model object after tensor decomposition. The data type is torch.nn.Module.
Dictionary consisting of the layer names before and after tensor decomposition, for example, {'conv1': ['conv1.0', 'conv1.1'], 'conv2': ['conv2.0', 'conv2.1'],...}.

Restrictions

The input model must be an object of the torch.nn.Module type.
This API function decomposes only the convolution constructed by using torch.nn.Conv2d().
This API automatically decomposes the convolutional layers that meet the decomposition conditions. For details about the conditions, see Restrictions.

Example

from amct_pytorch.tensor_decompose import auto_decomposition
net = Net()                                                    # Build a model object.
net.load_state_dict(torch.load("src_path/weights.pth"))        # Load model weights.
net, changes = auto_decomposition(                             # Perform tensor decomposition.
    model=net,
    decompose_info_path="decomposed_path/decompose_info.json"
)

If training is involved, this API must be called before the model parameters are passed to the optimizer; if torch.nn.parallel.DistributedDataParallel (DDP) is used, this API must be called before the model parameters are passed to the DDP.
This API modifies the input model object in place. That is, the model object input by the user is changed after decomposition (exception: The input model is a torch.nn.Conv2d object). In this case, this API does not modify it. If decomposition occurs, the newly constructed torch.nn.Module object is returned.

Parent topic: Tensor Decomposition APIs