分解流程

接口调用流程

接口调用流程如图1所示，具体分解示例请参见获取更多样例>tensor_decompose。

图1 张量分解接口调用流程

1. 在线分解流程

   在训练脚本中，准备好含有预训练权重的torch.nn.Module模型对象，在将模型参数传递给优化器之前，将模型对象传递给auto_decomposition接口进行张量分解，得到分解后的模型对象，即可直接对其进行finetune。

2. 离线分解流程
   1. 在任意脚本中，准备好含有预训练权重的torch.nn.Module模型对象，将模型对象和分解信息文件保存路径传递给auto_decomposition接口进行张量分解，得到分解后的模型对象和保存的分解信息文件，并对分解后的模型权重进行保存。
   2. finetune时，在训练脚本中，在将模型参数传递给优化器之前，将模型对象和2.a中得到的分解信息文件路径传递给decompose_network接口，该接口会将模型结构修改为分解后的结构，然后加载2.a保存的分解后的模型权重，对模型进行finetune。
   

   离线分解时，2.a步骤在调用auto_decomposition接口后，用户需要自行保存分解后模型权重；2.b步骤在调用完成decompose_network后，用户需自行加载所保存的分解后的模型权重。

   此方案设计目的是方便用户自由控制权重文件的存取，例如在权重文件中存储自定义信息。

张量分解后会将1个卷积分解为2个串联的卷积，模型的其中一层卷积分解前后情况如图2所示。

图2 卷积分解前后示意图

调用示例

以下示例中，(*)表示用户已有的代码，...表示用户已有代码的省略，此处仅为示例，实际用户代码可能不同，请根据实际情况进行调整。

1. 在线张量分解
   在训练脚本中，调用auto_decomposition分解含有预训练权重的PyTorch模型，然后直接finetune。

   from amct_pytorch.tensor_decompose import auto_decomposition
   net = Net()                                          # (*) 构建模型对象
   net.load_state_dict(torch.load("src_path/net.pth"))  # (*) 加载模型权重
   net, changes = auto_decomposition(model=net)         # 执行张量分解
   optimizer = build_optimizer(net, ...)                # (*) 构建优化器（将模型参数传递给优化器）
   train(net, optimizer, ...)                           # (*) finetune

2. 离线张量分解
   1. 在任意脚本中，调用auto_decomposition分解含有预训练权重的PyTorch模型，保存分解信息文件和分解后的模型权重。

      from amct_pytorch.tensor_decompose import auto_decomposition
      net = Net()                                                                 # (*) 构建模型对象
      net.load_state_dict(torch.load("src_path/weights.pth"))                     # (*) 加载模型权重
      net, changes = auto_decomposition(                                          # 执行张量分解，并保存分解信息文件
          model=net,
          decompose_info_path="decomposed_path/decompose_info.json"               # 分解信息文件保存路径
      )
      torch.save(net.state_dict(), "decomposed_path/decomposed_weights.pth")      # 保存分解后的模型权重

   2. 在训练脚本中，调用decompose_network，根据2.a得到的分解信息文件将模型结构修改为分解后的结构，再加载2.a保存的分解后的模型权重，进行finetune。

      from amct_pytorch.tensor_decompose import decompose_network
      net = Net()                                                                 # (*) 构建用户模型对象
      net, changes = decompose_network(                                           # 加载分解信息文件，将模型结构修改为张量分解后的结构
          model=net,
          decompose_info_path="decomposed_path/decompose_info.json"               # 上一步保存的分解信息文件路径
      )
      net.load_state_dict(torch.load("decomposed_path/decomposed_weights.pth"))   # 加载上一步保存的分解后模型权重
      optimizer = build_optimizer(net, ...)                                       # (*) 构建优化器（将模型参数传递给优化器）
      train(net, optimizer, ...)                                                  # (*) finetune