模型迁移快速入门

- 使用MindIE Torch前,已参考《MindIE安装指南》中“安装开发环境”章节完成环境准备。
- 此快速入门以提前编译TorchScript路线为例介绍MindIE Torch模型迁移流程,torch.export和torch.compile路线可参见MindIE Torch编程范式,更多样例可参见样例参考。
- TorchScript模型导出详细指导可参见Pytorch官网指导,torch.export导出ExportedProgram详细指导可参见Pytorch官网指导。
MindIE Torch模型迁移流程分为:导入mindietorch框架、模型导出、模型编译、模型推理和资源释放五个步骤,使用示例如下所示:
# 1.导入mindietorch框架
import torch
import mindietorch
# 2.模型导出
class SimpleNet(torch.nn.Module): # 定义模型SimpleNet
def __init__(self):
super().__init__()
def forward(self, x, y):
return torch.add(x, y)
traced_model = torch.jit.trace(SimpleNet().eval(), (torch.ones([5, 5]), torch.ones([5, 5]))) # 模型TorchScript模型导出
# 3.模型编译
mindietorch.set_device(0) # 设置使用device 0设备
compiled_model = mindietorch.compile(traced_model, inputs=[mindietorch.Input((5, 5)), mindietorch.Input((5, 5))], soc_version="Ascendxxx") # 静态模型编译,Ascend310xxx为昇腾AI处理器类型,根据服务器设备类型配置
# 4.模型推理
x = torch.ones([5, 5])
y = torch.ones([5, 5])
result = compiled_model.forward(x.to("npu"), y.to("npu")) # 模型加速推理
print(result.to("cpu")) # 输出计算结果
# 5.资源释放
mindietorch.finalize()