2. 环境搭建:在 Jetson Nano/Orin 上安装 TensorRT

说实话,环境搭建这一步,看着简单,坑却不少。我见过太多人卡在版本兼容性上,一折腾就是半天。今天咱们就把 Jetson 上的 TensorRT 环境捋清楚。

2.1 先搞清楚你的 Jetson 型号

不同的 Jetson 设备,对应的 JetPack 版本不一样。JetPack 里包含了 CUDA、cuDNN、TensorRT 这些核心组件。我个人习惯是先查清楚板子的型号,再决定装哪个版本的 JetPack。

Jetson 型号 推荐 JetPack 版本 CUDA 版本 TensorRT 版本
Jetson Nano (4GB) JetPack 4.6.x CUDA 10.2 TensorRT 8.0.x
Jetson Xavier NX JetPack 5.0.x CUDA 11.4 TensorRT 8.4.x
Jetson Orin NX/Orin Nano JetPack 5.1.x 或 6.0 CUDA 11.4 / 12.2 TensorRT 8.5.x / 8.6.x
注意: 千万别自己手动去 NVIDIA 官网下载 CUDA 安装包。Jetson 上的 CUDA 是定制版,必须通过 JetPack 刷机或 apt 源安装。我曾经见过有人硬装 x86 的 CUDA,结果系统直接崩了。

2.2 刷机还是 apt?两种方式选哪个

搭建环境有两种主流方式:

  • SDK Manager 刷机:适合第一次用 Jetson 的新手。把整个系统镜像写到 SD 卡或 NVMe 上,所有组件一步到位。
  • apt 在线安装:适合已经跑起来的系统。直接加 NVIDIA 的 apt 源,按需安装组件。

我个人更推荐第二种。为什么呢?因为刷机太费时间了,而且每次都要重新配置系统。用 apt 安装,你可以只装需要的组件,灵活很多。

2.3 用 apt 安装 TensorRT 和依赖

假设你的 Jetson 已经刷好了基础系统。接下来就是加源、安装、验证三步走。

2.3.1 添加 NVIDIA apt 源

# 添加 NVIDIA 的 apt 仓库
sudo apt update
sudo apt install -y nvidia-jetpack

# 或者只安装核心组件(推荐)
sudo apt install -y nvidia-jetpack-runtime
sudo apt install -y nvidia-jetpack-dev
小技巧: 如果你空间有限,只装 nvidia-jetpack-runtime 和 nvidia-jetpack-dev 就够了。别装全量的 nvidia-jetpack,那会把所有示例代码和文档都塞进去,占好几个 G。

2.3.2 验证 CUDA 和 cuDNN

装完之后,先别急着跑模型。先确认 CUDA 和 cuDNN 能不能正常工作。

# 检查 CUDA 版本
nvcc --version

# 检查 cuDNN 版本
cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

# 跑一个 CUDA 示例
cd /usr/local/cuda/samples/1_Utilities/deviceQuery
sudo make
./deviceQuery

如果 deviceQuery 能正确显示你的 GPU 信息,说明 CUDA 环境没问题。嗯,这一步很重要,我遇到过有人跳过了验证,结果后面跑模型时才发现 CUDA 没装好。

2.3.3 验证 TensorRT

TensorRT 的验证更简单。直接查版本号,再跑一个内置的示例。

# 查看 TensorRT 版本
dpkg -l | grep tensorrt

# 跑一个 TensorRT 示例
cd /usr/src/tensorrt/samples/sampleMNIST
make
./sampleMNIST
关键点: TensorRT 的 Python 包是 tensorrt,不是 nvidia-tensorrt。安装命令是 pip3 install tensorrt。但要注意,Jetson 上的 TensorRT Python 包已经预装在系统里了,你只需要 import tensorrt 就能用。

2.4 配置 Python 开发环境

Python 环境我建议用虚拟环境管理。别直接在系统 Python 里装包,容易搞乱依赖。

# 安装 pip 和虚拟环境工具
sudo apt install -y python3-pip python3-venv

# 创建虚拟环境
python3 -m venv trt_env

# 激活环境
source trt_env/bin/activate

# 安装常用包
pip install numpy onnx onnxruntime
pip install tensorrt  # Jetson 上这个包已经预装,但可以升级
避坑指南: 我曾经在 Jetson Nano 上装 torch,直接用 pip install torch,结果装了个 x86 版本,跑不起来。后来才发现 Jetson 上的 PyTorch 需要从 NVIDIA 的官方 whl 包安装。记住,Jetson 是 ARM 架构,很多包没有预编译版本,得用 pip install --index-url https://developer.download.nvidia.com/compute/redist 来装。

2.5 配置 C++ 开发环境

C++ 开发环境相对简单。Jetson 系统里已经自带了 g++ 和 cmake。你只需要确认一下版本。

# 检查编译器
g++ --version
cmake --version

# 编译一个 TensorRT 示例
cd /usr/src/tensorrt/samples/sampleOnnxMNIST
make
./sampleOnnxMNIST

如果你需要自己写 CMakeLists.txt,记得链接 TensorRT 的库:

cmake_minimum_required(VERSION 3.10)
project(trt_demo)

find_package(CUDA REQUIRED)
find_library(TENSORRT_LIB nvinfer)
find_library(TENSORRT_PLUGIN nvinfer_plugin)

add_executable(trt_demo main.cpp)
target_link_libraries(trt_demo ${CUDA_LIBRARIES} ${TENSORRT_LIB} ${TENSORRT_PLUGIN})

2.6 环境验证全流程

最后,我建议你写一个简单的脚本,把整个环境串起来验证一遍。这样以后换设备或重装系统时,直接跑一遍脚本就能确认环境是否正常。

#!/bin/bash
# env_check.sh - 验证 TensorRT 环境

echo "=== CUDA 版本 ==="
nvcc --version

echo "=== cuDNN 版本 ==="
cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

echo "=== TensorRT 版本 ==="
dpkg -l | grep tensorrt | head -n 1

echo "=== Python TensorRT ==="
python3 -c "import tensorrt; print(tensorrt.__version__)"

echo "=== 设备信息 ==="
/usr/local/cuda/samples/1_Utilities/deviceQuery/deviceQuery | head -n 20
我的习惯: 把这个脚本放在 /opt/trt_env/ 目录下,每次换设备或重装系统后,先跑一遍。如果所有检查都通过,再开始部署模型。这能省下不少排查问题的时间。

2.7 环境搭建知识体系

下面这张图总结了整个环境搭建的核心流程和组件关系。你可以把它当作一个检查清单。

Jetson TensorRT 环境搭建流程 Jetson 硬件 JetPack 系统镜像 CUDA cuDNN TensorRT Python 环境 C++ 环境 验证脚本 运行示例:sampleMNIST 核心原则:版本匹配 → 组件验证 → 示例运行,三步缺一不可

环境搭建说难不难,说简单也不简单。关键是版本要匹配,验证要到位。按照上面的步骤走一遍,你的 Jetson 就能跑 TensorRT 了。接下来就可以开始真正的模型优化和部署了。