产品支持情况
产品 |
是否支持 |
|---|---|
√ |
|
√ |
|
√ |
|
√ |
功能说明
对用户输入的PyTorch模型对象进行张量分解，得到分解后的模型对象和分解前后层的对应名称，并保存分解信息文件（可选）。
函数原型
1 | model, changes = auto_decomposition(model, decompose_info_path=None) |
参数说明
参数名 |
输入/输出 |
使用限制 |
|---|---|---|
model |
输入 |
含义：待分解的含有预训练权重的PyTorch模型对象。在调用该接口时建议将模型放置于CPU而不是GPU上，以防分解时显存不足。 数据类型：torch.nn.Module |
decompose_info_path |
输入 |
含义：分解信息文件保存路径。将以json格式存储，因此建议使用.json扩展名。为None时不保存分解信息文件（默认）。 数据类型：string 默认值：None |
返回值说明
- 张量分解后的模型对象，类型为torch.nn.Module。
- 张量分解前后的对应层名构成的字典，类型为dict，形如{'conv1': ['conv1.0', 'conv1.1'], 'conv2': ['conv2.0', 'conv2.1'], ...}。
约束说明
- 用户输入的模型需为torch.nn.Module类型的对象。
- 本接口函数仅对通过torch.nn.Conv2d()构建的卷积进行分解。
- 用户调用本接口函数，接口函数对符合分解条件的卷积层进行自动分解，约束请参见分解约束。
调用示例
1 2 3 4 5 6 7 | 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对象）。