Nvidia TensorRT RTX 执行提供程序

Nvidia TensorRT RTX 执行提供程序是消费级硬件 (RTX PC) 上 GPU 加速的首选执行提供程序。与专注于数据中心环境的旧版 TensorRT 执行提供程序相比，它更易于使用，并且比 CUDA EP 性能更佳。以下是使其比我们的旧版 TensorRT 执行提供程序更适合 RTX PC 的一些优点：

• 占用空间更小
• 模型编译/加载时间大大加快。
• 在跨多个 RTX GPU 使用缓存模型方面具有更好的可用性。

ONNX Runtime 中的 Nvidia TensorRT RTX 执行提供程序利用 NVIDIA 的 TensorRT RTX 深度学习推理引擎（TODO：一旦 TRT RTX 文档可用，更正链接）来加速 RTX GPU 上的 ONNX 模型。Microsoft 和 NVIDIA 密切合作，将 TensorRT RTX 执行提供程序与 ONNX Runtime 集成。

目前 TensorRT RTX 支持 Ampere 或更新架构的 RTX GPU。对 Turing GPU 的支持即将推出。

目录

• 安装
• 从源代码构建
• 要求
• 用法
• 配置
• NV TensorRT RTX EP 缓存
• 性能调优

安装

请选择 Onnx Runtime 的 Nvidia TensorRT RTX 版本：https://runtime.onnx.org.cn/docs/install。（TODO！）

从源代码构建

请参阅 构建说明。

要求

ONNX Runtime	TensorRT-RTX	CUDA
主分支	1.0	12.0-12.9
1.22	1.0	12.0-12.9

用法

C/C++

const auto& api = Ort::GetApi();
Ort::SessionOptions session_options;
api.SessionOptionsAppendExecutionProvider(session_options, "NvTensorRtRtx", nullptr, nullptr, 0);
Ort::Session session(env, model_path, session_options);

C API 详情请见 此处。

Python

要使用 TensorRT RTX 执行提供程序，您必须在实例化 InferenceSession 时显式注册 TensorRT RTX 执行提供程序。

import onnxruntime as ort
sess = ort.InferenceSession('model.onnx', providers=['NvTensorRtRtxExecutionProvider'])

配置

TensorRT RTX 设置可以通过 TensorRT 执行提供程序会话选项进行配置。

以下是设置 NV TensorRT RTX EP 会话选项的示例和不同场景

点击下方查看 Python API 示例

import onnxruntime as ort

model_path = '<path to model>'

# note: for bool type options in python API, set them as False/True
provider_options = {
  'device_id': 0,
  'nv_dump_subgraphs': False,
  'nv_detailed_build_log': True,
  'user_compute_stream': stream_handle
}

sess_opt = ort.SessionOptions()
sess = ort.InferenceSession(model_path, sess_options=sess_opt, providers=[('NvTensorRTRTXExecutionProvider', provider_options)])

点击下方查看 C++ API 示例

Ort::SessionOptions session_options;

cudaStream_t cuda_stream;
cudaStreamCreate(&cuda_stream);

// Need to put the CUDA stream handle in a string
char streamHandle[32];
sprintf_s(streamHandle, "%lld", (uint64_t)cuda_stream);

const auto& api = Ort::GetApi();
std::vector<const char*> option_keys = {
    "device_id",
    "user_compute_stream",  // this implicitly sets "has_user_compute_stream"
};
std::vector<const char*> option_values = {
    "1",
    streamHandle
};

Ort::ThrowOnError(api.SessionOptionsAppendExecutionProvider(session_options, "NvTensorRtRtx", option_keys.data(), option_values.data(), option_keys.size()));

场景

场景	NV TensorRT RTX EP 会话选项	类型
指定用于执行的 GPU ID	device_id	int
为 GPU 操作设置自定义计算流	user_compute_stream	string
设置 TensorRT RTX EP GPU 内存使用限制	nv_max_workspace_size	int
转储优化后的子图以进行调试	nv_dump_subgraphs	bool
捕获 CUDA 图以减少启动开销	nv_cuda_graph_enable	bool
启用构建步骤的详细日志记录	nv_detailed_build_log	bool
定义最小形状	nv_profile_min_shapes	string
定义最大形状	nv_profile_max_shapes	string
定义最佳形状	nv_profile_opt_shapes	string

注意：对于布尔类型选项，在 Python 中赋值为 True/False，在 C++ 中赋值为 1/0。

执行提供程序选项

TensorRT RTX 配置可以通过执行提供程序选项设置。当每个模型和推理会话都有自己的配置时，这会很有用。所有配置都应显式设置，否则将采用默认值。

device_id

• 描述：GPU 设备 ID。
• 默认值：0

user_compute_stream

• 描述：定义推理运行的计算流。它隐式设置了 has_user_compute_stream 选项。流句柄需要以十进制数字的形式打印成字符串，并作为会话选项传递，如上例所示。

• 这也可以通过 Python API 进行设置。

  • 例如，从 PyTorch 捕获的 cuda 流可以传递给 ORT-NV TensorRT RTX EP。点击下方查看示例代码

    import onnxruntime as ort
    import torch
    ...
    sess = ort.InferenceSession('model.onnx')
    if torch.cuda.is_available():
        s = torch.cuda.Stream()
        provider_options = {
          'device_id': 0,
          'user_compute_stream': str(s.cuda_stream)
        }
    
        sess = ort.InferenceSession(
          model_path,
          providers=[('NvTensorRtRtxExecutionProvider', provider_options)]
        )
    
        options = sess.get_provider_options()
        assert "NvTensorRtRtxExecutionProvider" in options
        assert options["NvTensorRtRtxExecutionProvider"].get("user_compute_stream", "") == str(s.cuda_stream)
    ...
    

• 为了利用用户计算流，建议使用 I/O 绑定将输入和输出绑定到设备中的张量。

nv_max_workspace_size

• 描述：TensorRT RTX 引擎的最大工作空间大小（以字节为单位）。

• 默认值：0（让 TensorRT 选择最佳值）。

nv_dump_subgraphs

• 描述：如果 ONNX 被拆分为多个执行提供程序，则转储子图。
  • 这有助于调试子图，例如使用 trtexec --onnx subgraph_1.onnx 并检查解析器的输出。

nv_detailed_build_log

• 描述：在 NV TensorRT RTX EP 上启用详细的构建步骤日志记录，并包含每个引擎构建的时间。

nv_cuda_graph_enable

• 描述：这将捕获一个 CUDA 图，这对于包含许多小层的网络来说可以显著提高性能，因为它减少了 CPU 上的启动开销。

nv_profile_min_shapes

nv_profile_max_shapes

nv_profile_opt_shapes

• 描述：使用包含所提供最小/最大/最佳形状的配置文件构建显式动态形状。
  • 默认情况下，TensorRT RTX 引擎将支持动态形状，为了提高性能，可以指定一个或多个显式形状范围。
  • 配置文件形状的格式为 input_tensor_1:dim_1xdim_2x...,input_tensor_2:dim_3xdim_4x...,...
    • 这三个标志都应提供才能启用显式配置文件形状功能。
  • 请注意，可以通过为同一输入张量传递多个形状来启用多个 TensorRT RTX 配置文件。
  • 有关更多详细信息，请查看 TensorRT 文档 优化配置文件。

NV TensorRT RTX EP 缓存

主要有两种 TRT RTX EP 缓存：

• 嵌入式引擎模型 / EPContext 模型
• 内部 TensorRT RTX 缓存

内部 TensorRT RTX 缓存由 EP 自动管理。用户只需要管理 EPContext 缓存。缓存对于大幅缩短会话创建时间至关重要。

TensorRT RTX 将编译分为预编译（AOT）引擎和即时（JIT）编译。AOT 编译可以存储为 EPContext 模型，此模型在多个 GPU 代次之间兼容。加载此类 EPContext 模型后，TensorRT RTX 将即时编译引擎以适应所使用的 GPU。此 JIT 过程通过 TensorRT RTX 的内部缓存加速。有关示例用法，请参见：https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc

更多关于嵌入式引擎模型 / EPContext 模型

• TODO：决定默认使用去权重引擎的计划。修复 EP 实现以启用该功能。解释动机并在此文档中提供如何使用正确选项的示例。
• EPContext 模型还支持封装外部编译的引擎，例如使用 trtexec。Python 工具中包含一个能够将此类预编译引擎封装到 ONNX 文件中的Python 脚本。（TODO：文档说明这如何与去权重引擎一起工作）。

性能调优

有关性能调优，请参阅此页面上的指南：ONNX Runtime 性能调优

当/如果使用 onnxruntime_perf_test 时，请使用标志 -e nvtensorrttrx。

TensorRT RTX 插件支持

TensorRT RTX 不支持插件。