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