快捷方式

安装

预编译二进制文件

Torch-TensorRT 2.x 主要围绕 Python 构建。因此,可以在 pypi.org 上找到预编译版本

依赖项

您需要安装 CUDA、PyTorch 和 TensorRT(Python 包即可)才能使用 Torch-TensorRT

安装 Torch-TensorRT

您可以使用以下命令安装 Python 包

python -m pip install torch torch-tensorrt tensorrt

已上传适用于 x86 Linux 和 Windows 的软件包

为特定 CUDA 版本安装 Torch-TensorRT

与 PyTorch 类似,Torch-TensorRT 具有针对不同 CUDA 版本编译的构建。这些构建发布在 PyTorch 的软件包索引上

例如 CUDA 11.8

python -m pip install torch torch-tensorrt tensorrt --extra-index-url https://download.pytorch.org/whl/cu118

安装夜间构建版本

Torch-TensorRT 发布了针对 PyTorch 夜间构建版本的夜间构建版本。可以从 PyTorch 夜间构建版本的软件包索引安装这些版本(按 CUDA 版本分开)

python -m pip install --pre torch torch-tensorrt tensorrt --extra-index-url https://download.pytorch.org/whl/nightly/cu128

C++ 预编译二进制文件(仅限 TorchScript)

发布版本的预编译 tar 包在此处提供:https://github.com/pytorch/TensorRT/releases

从源编译

在 Linux 上构建

依赖项

  • Torch-TensorRT 使用 Bazel 构建,因此首先安装 Bazel。

    export BAZEL_VERSION=$(cat <PATH_TO_TORCHTRT_ROOT>/.bazelversion)
    mkdir bazel
    cd bazel
    curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-dist.zip
    unzip bazel-$BAZEL_VERSION-dist.zip
    bash ./compile.sh
    cp output/bazel /usr/local/bin/
    
  • 您还需要在系统上安装 CUDA(如果在容器中运行,则系统必须安装 CUDA 驱动程序,容器必须安装 CUDA)

  • bazel 将为您拉取正确的 LibTorchTensorRT 版本。

    注意:默认情况下,bazel 将从 pytorch.org 拉取最新的夜间构建版本。对于构建主分支来说,这通常足够,但是如果您的目标是特定的 PyTorch 版本,请使用更新的 URL/路径编辑这些位置

  • 系统上不需要安装 TensorRT 即可构建 Torch-TensorRT,实际上这样更利于确保构建的可重现性。如果需要非默认版本,请将 WORKSPACE 文件指向 tarball 的 URL,或者从 https://developer.nvidia.com 下载 TensorRT 的 tarball 并在此处更新 WORKSPACE 文件中的路径 https://github.com/pytorch/TensorRT/blob/4e5b0f6e860910eb510fa70a76ee3eb9825e7a4d/WORKSPACE#L71

    例如

    http_archive(
        name = "tensorrt",
        build_file = "@//third_party/tensorrt/archive:BUILD",
        sha256 = "<TENSORRT SHA256>", # Optional but recommended
        strip_prefix = "TensorRT-<TENSORRT VERSION>",
        urls = [
            "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/<TENSORRT DOWNLOAD PATH>",
            # OR
            "file:///<ABSOLUTE PATH TO FILE>/TensorRT-<TENSORRT VERSION>.Linux.x86_64-gnu.cuda-<CUDA VERSION>.tar.gz"
        ],
    )
    

    请记住,在运行时,必须将这些库显式添加到您的 LD_LIBRARY_PATH

如果您安装了本地版本的 TensorRT,也可以使用它,方法是注释掉上面的行并取消注释下面的行 https://github.com/pytorch/TensorRT/blob/4e5b0f6e860910eb510fa70a76ee3eb9825e7a4d/WORKSPACE#L114C1-L124C3

构建软件包

正确配置 WORKSPACE 后,构建 torch-tensorrt 所需的命令如下

python -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu128

如果您使用 uv (https://docs.astral.sh/uv/) 工具管理 python 和您的项目,命令会稍微简单一些

uv pip install -e .

构建 wheel 文件

python -m pip wheel --no-deps --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu128 -w dist

其他构建选项

库中的某些功能是可选的,可以使构建更轻量或更具可移植性。

仅限 Python 分发版

库中有多个功能需要启用 C++ 组件。这包括接受 TorchScript 模块进行编译的 TorchScript 前端,以及 Torch-TensorRT 运行时(它是使用 Torch-TensorRT 编译的模块的默认执行器,无论是使用 TorchScript 还是 Dynamo 前端)。

如果您想要不需要 C++ 的构建,可以禁用这些功能并避免构建这些组件。因此,唯一可用的运行时将是基于 Python 的运行时,这会影响序列化等功能。

PYTHON_ONLY=1 python -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu128
不包含 TorchScript 前端

TorchScript 前端是 Torch-TensorRT 的一个旧版功能,现在处于维护状态,因为 TorchDynamo 已成为此项目首选的编译器技术。它包含相当多的 C++ 代码,对大多数用户来说不再必要。因此,您可以从构建中排除此组件以加快构建时间。基于 C++ 的运行时仍然可用。

NO_TORCHSCRIPT=1 python -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu128

单独构建 C++ 库(仅限 TorchScript)

Release 版本构建
bazel build //:libtorchtrt -c opt

包含头文件和库的 tarball 可以在 bazel-bin 中找到

Debug 版本构建

要构建包含调试符号的版本,请使用以下命令

bazel build //:libtorchtrt -c dbg

包含头文件和库的 tarball 可以在 bazel-bin 中找到

Pre CXX11 ABI 版本构建

要使用 pre-CXX11 ABI 构建,请使用 pre_cxx11_abi 配置

bazel build //:libtorchtrt --config pre_cxx11_abi -c [dbg/opt]

包含头文件和库的 tarball 可以在 bazel-bin 中找到

选择正确的 ABI

可能 Torch-TensorRT 编译中最复杂的事情是选择正确的 ABI。有两个相互不兼容的选项:pre-cxx11-abi 和 cxx11-abi。复杂性在于,虽然最流行的 PyTorch 分发版(直接从 pytorch.org/pypi 下载的 wheel 文件)使用 pre-cxx11-abi,但您可能会遇到的大多数其他分发版(例如来自 NVIDIA 的 - NGC 容器、适用于 Jetson 的构建以及某些 libtorch 构建,如果您从源构建 PyTorch,也可能如此)使用 cxx11-abi。重要的是您使用正确的 ABI 编译 Torch-TensorRT 以使其正常工作。下面是 PyTorch 分发源和推荐命令的一般配对表

PyTorch 来源

推荐的 Python 编译命令

推荐的 C++ 编译命令

来自 PyTorch.org 的 PyTorch whl 文件

python -m pip install .

bazel build //:libtorchtrt -c opt --config pre_cxx11_abi

来自 PyTorch.org 的 libtorch-shared-with-deps-*.zip

python -m pip install .

bazel build //:libtorchtrt -c opt --config pre_cxx11_abi

来自 PyTorch.org 的 libtorch-cxx11-abi-shared-with-deps-*.zip

python setup.py bdist_wheel

bazel build //:libtorchtrt -c opt

预装在 NGC 容器中的 PyTorch

python setup.py bdist_wheel

bazel build //:libtorchtrt -c opt

来自 NVIDIA Jetson 论坛的 PyTorch

python setup.py bdist_wheel

bazel build //:libtorchtrt -c opt

从源构建的 PyTorch

python setup.py bdist_wheel

bazel build //:libtorchtrt -c opt

注意:对于所有上述情况,您必须在 WORKSPACE 文件中正确声明您打算在 Python 和 C++ 构建中使用的 PyTorch 来源。请参见下方了解更多信息

在 Windows 上构建

  • Microsoft VS 2022 工具

  • Bazelisk

  • CUDA

构建步骤

  • 打开应用程序“适用于 VS 2022 的 x64 本机工具命令提示符” - 请注意,可能需要管理员权限

  • 确保 Bazelisk (Bazel 启动器) 已安装在您的机器上并可通过命令行访问。可以使用 Chocolatey 等软件包安装程序安装 Bazelisk

  • 安装最新版本的 Torch(例如,使用 pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 命令)

  • 克隆 Torch-TensorRT 仓库并导航到其根目录

  • 运行 pip install ninja wheel setuptools

  • 运行 pip install --pre -r py/requirements.txt

  • 运行 set DISTUTILS_USE_SDK=1

  • 运行 python setup.py bdist_wheel

  • 运行 pip install dist/*.whl

高级设置和故障排除

WORKSPACE 文件中,cuda_winlibtorch_wintensorrt_win 是特定于 Windows 的模块,可以进行自定义。例如,如果您想使用不同版本的 CUDA 进行构建,或者您的 CUDA 安装位于非标准位置,请更新 cuda_win 模块中的 path。

同样,如果您想使用不同版本的 pytorch 或 tensorrt,请分别自定义 libtorch_wintensorrt_win 模块中的 urls

这些软件包的本地版本也可以在 Windows 上使用。有关在 Windows 上使用本地版本 TensorRT 的示例,请参阅 toolchains\\ci_workspaces\\WORKSPACE.win.release.tmpl

其他构建系统

使用 CMake 构建(仅限 TorchScript)

可以使用 CMake 而不是 Bazel 构建 API 库 (位于 cpp/) 和 torchtrtc 可执行文件。目前,python API 和测试无法使用 CMake 构建。首先安装 CMake。

  • 最新版本的 CMake 和安装说明可在 [其网站](https://cmake.com.cn/download/) 上找到,适用于不同平台。

一些有用的 CMake 选项包括

  • TensorRT 的 CMake finder 位于 cmake/Modules 中。为了让 CMake 使用它们,在使用 CMake 配置项目时传递 -DCMAKE_MODULE_PATH=cmake/Modules

  • Libtorch 提供了自己的 CMake finder。如果 CMake 找不到它,请使用 -DTorch_DIR=/share/cmake/Torch 传递您的 libtorch 安装路径

  • 如果使用提供的 cmake finder 未找到 TensorRT,请指定 -DTensorRT_ROOT=

  • 最后,在您选择的构建目录中配置和构建项目,从 Torch-TensorRT 项目的根目录运行以下命令

cmake -S. -B<build directory> \
    [-DCMAKE_MODULE_PATH=cmake/Module] \
    [-DTorch_DIR=<path to libtorch>/share/cmake/Torch] \
    [-DTensorRT_ROOT=<path to TensorRT>] \
    [-DCMAKE_BUILD_TYPE=Debug|Release]
cmake --build <build directory>

在 aarch64 (Jetson) 上原生构建

前提条件

安装或编译适用于 aarch64 的 PyTorch/LibTorch 构建版本

NVIDIA 在此托管 Jetson 的最新发布分支构建版本

环境设置

要在 aarch64-linux-gnu 平台上原生构建,请使用本地可用依赖项配置 WORKSPACE

  1. WORKSPACE 替换为 //toolchains/jp_workspaces 中对应的 WORKSPACE 文件

  2. new_local_repository 规则中配置包含本地依赖项的目录根路径

    注意:如果您使用 pip 软件包安装了 PyTorch,则正确的路径是 python torch 软件包的根路径。如果您使用 sudo pip install 进行安装,路径将是 /usr/local/lib/python3.8/dist-packages/torch。如果您使用 pip install --user 进行安装,路径将是 $HOME/.local/lib/python3.8/site-packages/torch

如果您使用的是 NVIDIA 编译的 pip 软件包,请将两个 libtorch 源的路径设置为相同的路径。这是因为与 x86_64 上的 PyTorch 不同,NVIDIA aarch64 PyTorch 使用 CXX11-ABI。如果您使用 pre_cxx11_abi 从源编译,并且只想使用该库,请将路径设置为相同的路径,但在编译时确保添加 --config=pre_cxx11_abi 标志

new_local_repository(
    name = "libtorch",
    path = "/usr/local/lib/python3.8/dist-packages/torch",
    build_file = "third_party/libtorch/BUILD"
)

new_local_repository(
    name = "libtorch_pre_cxx11_abi",
    path = "/usr/local/lib/python3.8/dist-packages/torch",
    build_file = "third_party/libtorch/BUILD"
)
编译 C++ 库和编译器 CLI

注意:由于 Jetpack 4.5 和 4.6 之间依赖项位置的变化,现在有一个标志用于通知 bazel Jetpack 版本

--platforms //toolchains:jetpack_x.x

使用 bazel 命令编译 Torch-TensorRT 库

bazel build //:libtorchtrt --platforms //toolchains:jetpack_5.0
编译 Python API

注意:由于 Jetpack 4.5 和更新的 Jetpack 版本之间依赖项位置的变化,现在 setup.py 有一个用于设置 jetpack 版本的标志(默认:5.0)

//py 目录使用以下命令编译 Python API

python3 setup.py install

如果您的 PyTorch 构建版本使用 Pre-CXX11 ABI,请移除 --use-pre-cxx11-abi 标志

如果您正在为 Jetpack 4.5 构建,添加 --jetpack-version 5.0 标志

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源