MindIE Torch快速上手

MindIE Torch整体介绍

整体架构

推理迁移工作流

关键特性

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

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

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

cpp 复制代码
// 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路线):

python 复制代码
# 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模型进行编译优化
python 复制代码
# 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进行编译优化
python 复制代码
# 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进行即时编译优化。
python 复制代码
# 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
python 复制代码
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))] 
  • 动态分档
python 复制代码
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
python 复制代码
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开发

python 复制代码
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()
相关推荐
猿java几秒前
Elasticsearch有哪几种分页方式?该如何选择?
后端·elasticsearch·架构
绝无仅有5 分钟前
服务器Docker 安装和常用命令总结
后端·面试·github
ManageEngineITSM18 分钟前
云原生环境下的ITSM新趋势:从传统运维到智能化服务管理
大数据·运维·人工智能·云原生·itsm·工单系统
Chiy27 分钟前
架构设计避坑指南:读写分离后,分表分库到底该何时用?
后端
aneasystone本尊44 分钟前
可视化探索 GraphRAG 的知识图谱
人工智能
嘀咕博客1 小时前
Krea Video:Krea AI推出的AI视频生成工具
人工智能·音视频·ai工具
As33100101 小时前
Manus AI 与多语言手写识别技术全解析
大数据·网络·人工智能
张志鹏PHP全栈1 小时前
Rust第四天,Rust中常见编程概念
后端·rust
小璐乱撞1 小时前
超越传统 RAG:GraphRAG 全流程解析与实战指南
人工智能·后端
慧星云1 小时前
魔多 AI 上线提现功能 :将你的收益安稳入袋!
人工智能·云计算·aigc