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 将 Jetson 设备刷写为 JetPack 6.2：

   - 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/
   

构建 Torch-TensorRT

构建环境设置

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
   

构建 Wheel 文件

安装

安装后验证

通过在 Python 中导入来验证安装：.. code-block:: python

# 验证 torch-tensorrt 是否可以导入 import torch import torch_tensorrt print(torch_tensorrt.__version__)

# 验证示例是否可以运行 python examples/dynamo/torch_compile_resnet_example.py