昇腾社区首页
中文
注册

MindIE-Torch编程范式

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
    npuResults = 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
    npuResults = compiled_module.forward(input_data)
  • Python伪代码(适用于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
    npuResults = compiled_module.forward(input_data)