昇腾社区首页
中文
注册

模型迁移快速入门

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()