模型推理

推理输入输出

编译优化后的torchscript模型或nn.Module可调用torch的forward接口进行推理，forward的输入数据只支持"npu"的Tensor。模型的输出Tensor的设备与模型输入Tensor的设备有关，也与模型编译后是否是整图有关，如下表所示：

异步推理

MindIE-Torch模型推理采用异步推理，用户可以显式的设置当前Stream和进行Stream同步，也可以不显式设置，采用默认Stream。当用户不显式设置时，MindIE-Torch内部将会在数据拷贝时进行Stream同步，保证计算结果正确性。

• 不显式设置Stream的伪代码：
  • C++：

    dataNpu = dataCpu.to("npu:0");
    npuResults = compiled_module.forward(dataNpu); // 异步推理，采用默认流
    cpuResults = npuResults.to("cpu"); // torch_aie内部会在数据拷贝前进行流同步

  • Python：

    dataNpu = dataCpu.to("npu:0")
    npuResults = compiled_module.forward(dataNpu) # 异步推理，采用默认流
    cpuResults = npuResults.to("cpu") # MindIE-Torch内部会在数据拷贝前进行流同步

• 显式设置Stream的伪代码：
  • C++：

    c10::Stream stream = c10::Stream(c10::Stream::DEFAULT, npu); // create stream
    c10::StreamGuard streamGuard(stream); // set stream
    npuResults = compiledModule.forward(dataNpu); // forward
    stream.synchronize() // synchronize

  • Python：

    stream = mindietorch.npu.Stream("npu:0") # create stream
    with mindietorch.npu.stream(stream): # set stream
         npu_results = compiled_module.forward(data_npu) # forward
    stream.synchronize() # synchronize