MindIE Torch快速上手

MindIE Torch整体介绍

整体架构

推理迁移工作流

关键特性

特性1：配合torch_npu实现子图+单算子混合执行

特性2：支持C++和Python编程

c++伪代码(适用于TorchScript路线)：

// load TorchScript module
torch::jit::script::Module module = torch::jit::load("xxx.pth"); 

// step1:compile
torch_aie::torchscript::CompileSpec compileSpec(inputsInfo);
auto compiled_module = torch_aie::torchscript::compile(module, compileSpec);

 // step2:forward
npu_results = compiled_module.forward(input);

python伪代码(适用于TorchScript路线)：

# load TorchScript module
model = torch.jit.load("xxx.pth") 

# step1:compile
compiled_module = mindietorch.compile(model, inputs=inputs_info) 

# step2:forward
npu_results = compiled_module.forward(input_data)

特性3：支持TorchScript、ExportedProgram多种模式

• TorchScript模式
  支持对torch.jit.trace/script导出的TorchScript模型进行编译优化

# load TorchScript module
model = torch.jit.load("xxx.pth") 

# step1:compile
compiled_module = mindietorch.compile(model, inputs=inputs_info) 

# step2:forward
npu_results = compiled_module.forward(input_data)

• torch.export 模式
  支持对torch.export导出的ExportedProgram进行编译优化

# load PyTorch nn.module
model = torch.load("xxx.pth")
# 在使用MindIE Torch之前用户可以选择提前导出ExportedProgram
exported_model = torch.export.export(model, args=tuple(input_data,))

# step1:compile
# 当传入一个nn.Module时, MindIE Torch内部会先导出ExportedProgram, 再进行编译优化
compiled_module = mindietorch.compile(model, inputs=inputs_info, ir="dynamo") 
# 当传入一个ExportedProgram时, MindIE Torch会直接进行编译优化
compiled_module = mindietorch.compile(exported_model, inputs=inputs_info, ir="dynamo")

# step2:forward
npu_results = compiled_module.forward(input_data)

• torch.compile 路线
  提供了名为mindie的编译后端，支持在推理时对torch.compile;
  生成的GraphModule进行即时编译优化。

# load PyTorch nn.module
model = torch.load("xxx.pth")

# step1:准备待执行模型，此时并不会进行模型的编译优化
opt_model = torch.compile(model, backend="mindie")
# 或
opt_model = mindietorch.compile(model, ir="torch_compile")

# step2:执行推理，此时PyTorch才会调用MindIE Torch的自定义后端进行模型编译，编译完成之后才开始推理
npu_results = opt_model(input_data)

特性4：支持静态Shape、动态Shape模型编译

• 静态shape

import torch 
import mindietorch

# 纯静态shape，单输入
inputs = [mindietorch.Input((batchsize, 3, 224, 224))] # batchsize 自行写入

# 纯静态shape，多输入
inputs = [mindietorch.Input((batchsize, 3, 224, 224)), mindietorch.Input((batchsize, 3, 224, 224))] 

• 动态分档

import torch 
import mindietorch

# 动态分档
inputs = []
inputs_gear_1 = [mindietorch.Input((1, 3, 224, 224))]  # 1档
inputs_gear_2 = [mindietorch.Input((8, 3, 224, 224))]  # 2档
inputs_gear_3 = [mindietorch.Input((32, 3, 224, 224))]  # 3档
inputs.append(inputs_gear_1)
inputs.append(inputs_gear_2)
inputs.append(inputs_gear_3) 

• ShapeRange

import torch 
import mindietorch

# input shape range
min_shape = (1, 3, 224, 224)
max_shape = (32, 3, 224, 224)
inputs = []
inputs.append(mindietorch.Input(min_shape = min_shape, max_shape= max_shape)) # inputs 作为最终的模型输入设置。

MindIE Torch demo开发

import torch
# 导入mindietorch
import mindietorch 

# 加载原始ts模型、构造输入数据
model_path = "./resnet50.ts"
model = torch.jit.load(model_path)

# 1. 构造输入
inputs = [mindietorch.Input((batchsize, 3, 224, 224))]

# 2. 编译优化
mindietorch.set_device(0)
compiled_model = mindietorch.compile(model, inputs=inputs, ir="ts") 

# 3.模型推理
results = compiled_model.forward(input_data.to("npu"))

# 4. 编译后模型保存加载 (可选)
compiled_module.save("CompiledModel.ts")

# 5. om离线模型导出(可选)
mindietorch.export_engine(model, "forward", "export_engine.om", inputs=inputs)

# 6. 资源释放
mindietorch.finalize()