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 |
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,不是 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 上这个包已经预装,但可以升级
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 了。接下来就可以开始真正的模型优化和部署了。