auto_decomposition

功能说明

对用户输入的PyTorch模型对象进行张量分解，得到分解后的模型对象和分解前后层的对应名称，并保存分解信息文件（可选）。

约束说明

用户输入的模型需为torch.nn.Module类型的对象。
本接口函数仅对通过torch.nn.Conv2d()构建的卷积进行分解。
用户调用本接口函数，接口函数对符合分解条件的卷积层进行自动分解，约束请参见分解约束。

函数原型

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

参数说明

参数名	输入/返回值	含义	使用限制
model	输入	待分解的含有预训练权重的PyTorch模型对象。在调用该接口时建议将模型放置于CPU而不是GPU上，以防分解时显存不足。	数据类型：torch.nn.Module
decompose_info_path	输入（可选）	分解信息文件保存路径。将以json格式存储，因此建议使用.json扩展名。为None时不保存分解信息文件（默认）。	数据类型：string
model	返回值	张量分解后的模型对象。	数据类型：torch.nn.Module
changes	返回值	张量分解前后的对应层名构成的字典，形如{'conv1': ['conv1.0', 'conv1.1'], 'conv2': ['conv2.0', 'conv2.1'], ...}。	数据类型：dict

返回值说明

张量分解后的模型对象、张量分解前后的对应层名构成的字典。

函数输出

decompose_info_path：保存的分解信息文件，以json格式存储。

调用示例

      
           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.nn.parallel.DistributedDataParallel (DDP)，则本接口的调用也需在将模型传递给DDP之前。
本接口将原地修改传入的模型对象，即分解后会改变用户传入的模型对象本身（例外：传入的模型是一个torch.nn.Conv2d对象，该情况下本接口不会对其进行修改，如发生分解，则返回的模型是新构建的torch.nn.Module对象）。

父主题： 张量分解接口