提前编译模式下，MindIE Torch主要分为compile和forward两阶段工作。
- 第一阶段，用户调用compile接口对TorchScript模型/ExportedProgram进行编译优化，返回一个可以在NPU上加速推理的TorchScript模型/nn.Module。
- 第二阶段，用户准备模型输入数据，调用torch的模型forward接口，即可完成模型NPU推理。
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)
|
- 适用于torch.export路线：
1
2
3
4
5
6
7
8
9
10
11
12
13 | # 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)
|