Torch-TensorRT-RTX¶
概述¶
TensorRT-RTX¶
TensorRT for RTX 基于 NVIDIA TensorRT 推理库久经考验的性能，并简化了 AI 模型在 NVIDIA RTX GPU（跨台式机、笔记本电脑和工作站）上的部署。
TensorRT for RTX 是 NVIDIA TensorRT 的即插即用替代品，适用于面向从 Turing 到 Blackwell 代的 NVIDIA RTX GPU 的应用程序。它在运行时引入了一个即时 (JIT) 优化器，该优化器可以在最终用户的 RTX 加速 PC 上直接编译改进的推理引擎，耗时不到 30 秒。这消除了冗长的预编译步骤的需要，并实现了快速的引擎生成、改进的应用程序可移植性和前沿的推理性能。
有关 TensorRT-RTX 的详细信息，请参阅
目前，Torch-TensorRT 仅支持 TensorRT-RTX 用于实验目的。Torch-TensorRT 在构建和运行时默认使用标准 TensorRT。
要使用 TensorRT-RTX
使用
--use-rtx标志构建 wheel 包，或将USE_TRT_RTX=true设置为 true。运行时，将
USE_TRT_RTX=true环境变量设置为 true 以调用 TensorRT-RTX。
先决条件¶
安装 Bazel¶
使用 TensorRT-RTX 构建 wheel 包需要 Bazel。
在 Linux 上
curl -L https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-amd64 \
-o bazelisk \
&& mv bazelisk /usr/bin/bazel \
&& chmod +x /usr/bin/bazel
在 Windows 上
choco install bazelisk -y
安装 TensorRT-RTX Tarball¶
TensorRT-RTX tarball 可从 https://developer.nvidia.com/tensorrt-rtx 下载。目前，Torch-TensorRT 使用 TensorRT-RTX 版本 1.0.0.21。
下载后
在 Linux 上
请确保将 lib 路径添加到 LD_LIBRARY_PATH 环境变量。
# If TensorRT-RTX is downloaded in /your_local_download_path/TensorRT-RTX-1.0.0.21
export LD_LIBRARY_PATH=/your_local_download_path/TensorRT-RTX-1.0.0.21/lib:$LD_LIBRARY_PATH
echo $LD_LIBRARY_PATH | grep TensorRT-RTX
在 Windows 上
请确保将 lib 路径添加到 Windows 系统变量 PATH。
# If TensorRT-RTX is downloaded in C:\your_local_download_path\TensorRT-RTX-1.0.0.21
set PATH="%PATH%;C:\your_local_download_path\TensorRT-RTX-1.0.0.21\lib"
echo %PATH% | findstr TensorRT-RTX
安装 TensorRT-RTX Wheel¶
目前，tensorrt_rtx wheel 包尚未在 PyPI 上发布。您必须从下载的 tarball 中手动安装它。
# If the tarball is downloaded in /your_local_download_path/TensorRT-RTX-1.0.0.21
python -m pip install /your_local_download_path/TensorRT-RTX-1.0.0.21/python/tensorrt_rtx-1.0.0.21-cp39-none-linux_x86_64.whl
使用 TensorRT-RTX 构建 Torch-TensorRT¶
本地使用 TensorRT-RTX 构建¶
# If you have previously built with standard TensorRT, make sure to clean the build environment,
# otherwise it will use the existing .so built with standard TensorRT, which is not compatible with TensorRT-RTX.
python setup.py clean
bazel clean --expunge
#remove everything under build directory,
rm -rf build/*
# Build wheel with TensorRT-RTX
python setup.py bdist_wheel --use-rtx
# Install the wheel
python -m pip install dist/torch-tensorrt-*.whl
快速开始¶
# You must set USE_TRT_RTX=true to use TensorRT-RTX
USE_TRT_RTX=true python examples/dynamo/torch_compile_resnet_example.py
故障排除¶
如果遇到加载或链接错误，请检查 tensorrt_rtx 是否已正确链接。如果未正确链接，请清理环境并重新构建。
在 Linux 上
# Ensure only tensorrt_rtx is installed (no standard tensorrt wheels)
python -m pip list | grep tensorrt
# Check if libtorchtrt.so links to the correct tensorrt_rtx shared object
trt_install_path=$(python -m pip show torch-tensorrt | grep "Location" | awk '{print $2}')/torch_tensorrt
# Verify libtensorrt_rtx.so.1 is linked, and libnvinfer.so.10 is NOT
ldd $trt_install_path/lib/libtorchtrt.so
在 Windows 上
# Check if tensorrt_rtx_1_0.dll is linked, and libnvinfer.dll is NOT
cd py/torch_tensorrt
dumpbin /DEPENDENTS torchtrt.dll
