I/O 绑定

当使用非 CPU 执行提供程序时,最有效的方式是在执行图(调用 Run())之前,将输入(和/或输出)安排在目标设备上(由使用的执行提供程序抽象化)。当输入未复制到目标设备时,ORT 会在 Run() 调用的过程中从 CPU 复制输入。 类似地,如果输出未在设备上预先分配,则 ORT 会假定输出是在 CPU 上请求的,并在 Run() 调用的最后一步从设备复制输出。 这会占用图的执行时间,误导用户认为 ORT 很慢,而大部分时间都花费在这些复制上。

为了解决这个问题,我们引入了 IOBinding 的概念。 关键思想是在调用 Run() 之前,安排将输入复制到设备,并预先在设备上分配输出。 IOBinding 在我们所有的语言绑定中都可用。

以下是在各种语言中演示此功能用法的代码片段。

  • C++
      Ort::Env env;
      Ort::Session session(env, model_path, session_options);
      Ort::IoBinding io_binding{session};
      auto input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_tensor_values.data(), input_tensor_size, input_node_dims.data(), 4);
      io_binding.BindInput("input1", input_tensor);
      Ort::MemoryInfo output_mem_info{"Cuda", OrtDeviceAllocator, 0,
                                      OrtMemTypeDefault};
      // Use this to bind output to a device when the shape is not known in advance. If the shape is known you can use the other overload of this function that takes an Ort::Value as input (IoBinding::BindOutput(const char* name, const Value& value)).
      // This internally calls the BindOutputToDevice C API.
    
      io_binding.BindOutput("output1", output_mem_info);
      session.Run(run_options, io_binding);
    

    请注意,在上面的代码示例中,输出张量在绑定之前未分配,而是绑定了 Ort::MemoryInfo 作为输出。 这是一种有效的方式,让会话根据需要的形状分配张量。 特别是对于数据相关的形状或动态形状,这可能是一个很好的解决方案,可以获得正确的分配。 但是,如果输出形状已知且应重用输出张量,则最好将 Ort::Value 也绑定到输出。 这可以使用会话分配器或外部内存进行分配。 有关更多详细信息,请参阅设备张量文档

     Ort::Allocator gpu_allocator(session, output_mem_info);
     auto output_value = Ort::Value::CreateTensor(
          gpu_allocator, output_shape.data(), output_shape.size(),
          ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
     io_binding.BindOutput("output1", output_mem_info);
    
  • Python (请参阅 Python API 文档)

  • C# (请参阅 OrtIoBindingAllocationTest.cs)