Workflow

This section describes the API call sequence and example of tensor decomposition.

API Call Sequence

Figure 1 shows the API call sequence. For the decomposition example, see Sample List.

Figure 1 API call sequence for tensor decomposition

Online decomposition process
Prepare a torch.nn.Module model object that contains pre-trained weights. Before passing the model parameters to the optimizer, pass the model object to the auto_decomposition API for tensor decomposition. Then, directly fine-tune the decomposed model object.
Offline decomposition process
1. In any script, prepare a torch.nn.Module model object that contains pre-trained weights, and pass the model object and the decomposition information file save path to the auto_decomposition API to decompose the model object with the decomposition information file generated and save the decomposed weights.
2. During fine-tuning, in the training script, before passing the model parameters to the optimizer, pass the model object and the decomposition information file save path obtained in 2.a to the decompose_network API. This API then modifies the model structure based on the said file and loads the decomposed weights saved in 2.a for fine-tuning.
In offline decomposition, you need to save the decomposed weights yourself after calling the auto_decomposition API in 2.a and load the decomposed weights yourself after calling the decompose_network API in 2.b.

The purpose of this design is to facilitate the user to freely control the access to the weight file, for example, to store the custom information in the weight file.

After tensor decomposition, one convolution is decomposed into two concatenated convolutions. Figure 2 shows the convolutions before and after decomposition.

Figure 2 Diagrams before and after convolution decomposition

Example

In the following example, the asterisk (*) indicates the existing code, and the ellipsis (...) indicates the omitted code. The code is for reference only. The actual code may be different. Adjust the code based on your need.

Online tensor decomposition

In the training script, call the auto_decomposition API to decompose the PyTorch model with pre-trained weights, and then fine-tune the decomposed model.

         
              from amct_pytorch.tensor_decompose import auto_decomposition
net = Net()                                          # (*) Build a model object.
net.load_state_dict(torch.load("src_path/net.pth"))  # (*) Load model weights.
net, changes = auto_decomposition(model=net)         # Perform tensor decomposition.
optimizer = build_optimizer(net, ...)                # (*) Build an optimizer and pass model parameters to the optimizer.
train(net, optimizer, ...)                           # (*) fine-tune

Offline tensor decomposition

In any script, call the auto_decomposition API to decompose the PyTorch model with pre-trained weights and save the decomposition information file and the decomposed weights.

          
               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 and save the decomposition information file.
    model=net,
    decompose_info_path="decomposed_path/decompose_info.json"               # Set the decomposition information file save path.
)
torch.save(net.state_dict(), "decomposed_path/decomposed_weights.pth")      # Save the decomposed weights.

In the training script, call the decompose_network API to modify the model structure as described in the decomposition information file obtained in 2.a, and load the decomposed weights saved in 2.a for fine-tuning.

          
               from amct_pytorch.tensor_decompose import decompose_network
net = Net()                                                                 # (*) Build a model object.
net, changes = decompose_network(                                           # Load the decomposition information file to modify the model structure.
    model=net,
    decompose_info_path="decomposed_path/decompose_info.json"               # Specify the decomposition information file path set in the previous step.
)
net.load_state_dict(torch.load("decomposed_path/decomposed_weights.pth"))   # Load the decomposed weights saved in the previous step.
optimizer = build_optimizer(net, ...)                                       # (*) Build an optimizer and pass model parameters to the optimizer.
train(net, optimizer, ...)                                                  # (*) fine-tune

Parent topic: Tensor Decomposition