用户输入PyTorch模型对象和通过auto_decomposition保存的分解信息文件,根据分解信息文件将模型对象改变为张量分解后的结构,得到分解后的模型对象和分解前后层的对应名称。
model, changes = decompose_network(model, decompose_info_path)
参数名 |
输入/返回值 |
含义 |
使用限制 |
---|---|---|---|
model |
输入 |
待分解的PyTorch模型对象。在调用该接口时建议将模型放置于CPU而不是GPU上,以防分解时显存不足。 |
数据类型:torch.nn.Module |
decompose_info_path |
输入 |
分解信息文件路径,该文件通过auto_decomposition获得。 |
数据类型:string |
model |
返回值 |
改变为张量分解后结构的模型对象。 |
数据类型:torch.nn.Module |
changes |
返回值 |
张量分解前后的对应层名构成的字典,形如{'conv1': ['conv1.0', 'conv1.1'], 'conv2': ['conv2.0', 'conv2.1'], ...}。 |
数据类型:dict |
改变为张量分解后结构的模型对象、张量分解前后的对应层名构成的字典。
无。
1 2 3 4 5 6 |
from amct_pytorch.tensor_decompose import decompose_network net = Net() # 构建用户模型对象 net, changes = decompose_network( # 加载分解信息文件,将模型结构修改为张量分解后的结构 model=net, decompose_info_path="decomposed_path/decompose_info.json" # 由auto_decomposition保存的分解信息文件路径 ) |