安装¶
预编译二进制文件¶
Torch-TensorRT 2.x 主要以 Python 为中心。因此,预编译版本可以在 pypi.org 上找到
依赖项¶
您需要安装 CUDA、PyTorch 和 TensorRT(python 包就足够了)才能使用 Torch-TensorRT
安装 Torch-TensorRT¶
您可以使用以下命令安装 python 包
python -m pip install torch torch-tensorrt tensorrt
软件包已上传,适用于 Linux (x86) 和 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
安装 Nightly Builds¶
Torch-TensorRT 分发以 PyTorch nightly 为目标的 nightly 版本。这些版本可以从 PyTorch nightly 软件包索引(按 CUDA 版本分隔)安装
python -m pip install --pre torch torch-tensorrt tensorrt --extra-index-url https://download.pytorch.org/whl/nightly/cu124
C++ 预编译二进制文件(仅限 TorchScript)¶
发布的预编译 tarball 在此处提供:https://github.com/pytorch/TensorRT/releases
从源代码编译¶
在 Linux 上构建¶
依赖项¶
Torch-TensorRT 是使用 Bazel 构建的,因此首先安装它。
最简单的方法是使用您选择的方法安装 bazelisk https://github.com/bazelbuild/bazelisk
否则,您可以使用以下说明安装二进制文件 https://docs.bazel.build/versions/master/install.html
最后,如果您需要从源代码编译(例如,aarch64,直到 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)
如果不是正在构建的分支中使用的版本,请在此处指定您的 CUDA 版本:https://github.com/pytorch/TensorRT/blob/4e5b0f6e860910eb510fa70a76ee3eb9825e7a4d/WORKSPACE#L46
正确的 LibTorch 和 TensorRT 版本将由 bazel 为您拉取。
注意:默认情况下,bazel 将从 pytorch.org 拉取最新的 nightly 版本。对于构建 main 分支,这通常就足够了,但是,如果您有特定的 PyTorch 目标版本,请使用更新的 URL/路径编辑这些位置
构建 Torch-TensorRT 不需要在系统上安装 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/cu124
如果您使用 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/cu124 -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/cu124
无 TorchScript 前端¶
TorchScript 前端是 Torch-TensorRT 的一项遗留功能,现在处于维护状态,因为 TorchDynamo 已成为该项目的首选编译器技术。它包含相当多的 C++ 代码,对于大多数用户来说不再是必需的。因此,您可以从构建中排除此组件,以加快构建时间。基于 C++ 的运行时仍然可用。
NO_TORCHSCRIPT=1 python -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu124
独立构建 C++ 库(仅限 TorchScript)¶
发布版本构建¶
bazel build //:libtorchtrt -c opt
然后可以在 bazel-bin
中找到包含头文件和库的 tarball
调试版本构建¶
要使用调试符号构建,请使用以下命令
bazel build //:libtorchtrt -c dbg
然后可以在 bazel-bin
中找到包含头文件和库的 tarball
Pre CXX11 ABI 构建¶
要使用 pre-CXX11 ABI 构建,请使用 pre_cxx11_abi
配置
bazel build //:libtorchtrt --config pre_cxx11_abi -c [dbg/opt]
然后可以在 bazel-bin
中找到包含头文件和库的 tarball
选择正确的 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 文件中正确声明您打算使用的 PyTorch 来源,用于 Python 和 C++ 构建。有关更多信息,请参见下文
在 Windows 上构建¶
Microsoft VS 2022 工具
Bazelisk
CUDA
构建步骤¶
打开应用程序“x64 Native Tools Command Prompt for VS 2022” - 请注意,可能需要管理员权限
确保 Bazelisk(Bazel 启动器)已安装在您的计算机上,并且可以从命令行访问。可以使用 Chocolatey 等软件包安装程序来安装 Bazelisk
安装最新版本的 Torch(即使用
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124
)克隆 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_win
、libtorch_win
和 tensorrt_win
是特定于 Windows 的模块,可以自定义。例如,如果您想使用不同版本的 CUDA 构建,或者您的 CUDA 安装在非标准位置,请更新 cuda_win 模块中的 path。
同样,如果您想使用不同版本的 pytorch 或 tensorrt,请分别自定义 libtorch_win
和 tensorrt_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 查找器在 cmake/Modules 中提供。为了让 CMake 使用它们,请在使用 CMake 配置项目时传递 -DCMAKE_MODULE_PATH=cmake/Modules。
Libtorch 提供了自己的 CMake 查找器。如果 CMake 找不到它,请使用 -DTorch_DIR=<libtorch 路径>/share/cmake/Torch 传递 libtorch 安装路径
如果使用提供的 cmake 查找器找不到 TensorRT,请指定 -DTensorRT_ROOT=<TensorRT 路径>
最后,在您选择的构建目录中配置和构建项目,命令从 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
。
将
WORKSPACE
替换为//toolchains/jp_workspaces
中的相应 WORKSPACE 文件在
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
标志