昇腾社区首页
中文
注册

模型推理

推理输入输出

MindIE Torch编译优化后的模型在执行推理时,输入数据只支持"npu"的Tensor,输出数据也为"npu"的Tensor。

对于输入数据,"cpu"类型的Tensor需要to("npu")后进行推理运算;对于输出数据,"npu"类型的Tensor需要to("cpu")后再进行运算操作。

异步推理

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