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://onnxruntime.llms.tw/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 不支援外掛。