提前编译
提前编译模式下，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路线：
# 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)
父主题： MindIE Torch编程范式