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