JetPack 中的 Torch-TensorRT

概述

JetPack 6.2

NVIDIA JetPack 6.2 是 Jetson 平台的最新生产版本，包含： - CUDA 12.6 - TensorRT 10.3 - cuDNN 9.3

有关 JetPack 6.2 的详细信息，请参阅：* JetPack 6.2 发行说明 * Jetson 平台的 PyTorch

先决条件

系统准备

1. 使用 SDK Manager 刷写 JetPack 6.2 版本的 Jetson 设备

   : - SDK Manager 指南

2. 验证 JetPack 安装:

   apt show nvidia-jetpack
   

3. 安装开发组件： .. code-block:: sh

   sudo apt-get update sudo apt-get install nvidia-jetpack

4. 确认 CUDA 12.6 安装:

   nvcc --version
   # If missing or incorrect version:
   sudo apt-get install cuda-toolkit-12-6
   

5. 验证 cuSPARSELt 库:

   # Check library presence
   ls /usr/local/cuda/lib64/libcusparseLt.so
   
   # Install if missing
   wget https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-sbsa/libcusparse_lt-linux-sbsa-0.5.2.1-archive.tar.xz
   tar xf libcusparse_lt-linux-sbsa-0.5.2.1-archive.tar.xz
   sudo cp -a libcusparse_lt-linux-sbsa-0.5.2.1-archive/include/* /usr/local/cuda/include/
   sudo cp -a libcusparse_lt-linux-sbsa-0.5.2.1-archive/lib/* /usr/local/cuda/lib64/
   

JetPack 中的 Torch-TensorRT 安装

您可以直接从 JPL 存储库安装为 JetPack 6.2 构建的 torch-tensorrt wheel。

在 JetPack 中构建 Torch-TensorRT

您也可以在自己的设备上从源代码构建 torch-tensorrt wheel。

构建环境设置

1. 安装构建依赖项:

   wget https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-arm64
   sudo mv bazelisk-linux-arm64 /usr/bin/bazel
   sudo chmod +x /usr/bin/bazel
   

   apt-get install ninja-build vim libopenblas-dev git
   

2. 安装 Python 依赖项:

   wget https://bootstrap.pypa.io/get-pip.py
   python get-pip.py
   python -m pip install pyyaml
   

3. 安装 PyTorch:

   # Can only install the torch and torchvision wheel from the JPL repo which is built specifically for JetPack 6.2
   python -m pip install torch==2.8.0 torchvision==0.23.0  --index-url=https://pypi.jetson-ai-lab.io/jp6/cu126
   

4. 构建 Wheel:

   python setup.py bdist_wheel --jetpack
   

5. 安装 Wheel:

   # 您将在 dist 目录中找到 wheel cd dist python -m pip install torch_tensorrt-2.8.0.dev0+d8318d8fc-cp310-cp310-linux_aarch64.whl

6. 通过在 Python 中导入来验证安装:

# verify whether the torch-tensorrt can be imported
import torch
import torch_tensorrt
print(torch_tensorrt.__version__)

# verify whether the examples can be run
python examples/dynamo/torch_compile_resnet_example.py