4、环境搭建与工具链准备：PyTorch环境配置、ONNX与TensorRT安装、NVIDIA Jetson SDK Manager使用、RKNN Toolkit安装

说实话，很多同学在模型压缩和部署上栽跟头，不是因为算法不懂，而是环境没搭好。我见过太多人花了两周调代码，最后发现是CUDA版本对不上——这种坑，咱们得提前填平。

这一章，咱们就把所有工具链捋一遍。从PyTorch训练环境，到ONNX中间表示，再到TensorRT和RKNN这两个推理后端。嗯，一步到位。

4.1 PyTorch环境配置

PyTorch是整个流程的起点。我个人习惯用Miniconda管理环境，干净、隔离、好回滚。

创建环境的命令很简单：

conda create -n yolo11_deploy python=3.10
conda activate yolo11_deploy
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

装完之后，记得验证一下：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

如果输出 True，说明GPU能用。我曾经遇到一个学员，装完发现CUDA不可用，查了半天——原来是没装NVIDIA驱动。嗯，驱动是前提，别忘了。

4.2 ONNX与TensorRT安装

模型训练好之后，第一步是导出ONNX。ONNX像个中间语言，让模型能在不同框架间流转。

4.2.1 ONNX安装

pip install onnx onnxruntime-gpu

这里要注意，onnxruntime-gpu 和 onnxruntime 是两回事。前者支持CUDA加速，后者只有CPU。我刚开始做部署时就用错了，推理速度慢得离谱——后来才发现装的是CPU版本。

4.2.2 TensorRT安装

TensorRT是NVIDIA的推理优化引擎。它能把模型量化、融合算子、调整内存布局，说白了就是让模型跑得更快。

安装方式有两种：

方式	适用场景	命令/步骤
pip安装	x86_64 Linux桌面	pip install tensorrt==8.6.1
deb包安装	Jetson平台	通过JetPack SDK Manager安装
tar包安装	自定义路径/离线环境	下载后解压，配置LD_LIBRARY_PATH

验证安装：

python -c "import tensorrt as trt; print(trt.__version__)"

4.3 NVIDIA Jetson SDK Manager使用

Jetson是边缘端的主力硬件。Orin NX、Xavier NX、Nano，我都部署过。SDK Manager是刷机和装环境的官方工具。

使用流程：

1. 下载SDK Manager：从NVIDIA官网下载，需要注册账号
2. 连接Jetson：用USB线连接Jetson和宿主机，Jetson进入恢复模式
3. 选择JetPack版本：我建议用JetPack 5.1.3（对应CUDA 11.8），稳定
4. 选择组件：勾选CUDA、cuDNN、TensorRT、OpenCV
5. 开始刷机：等待30-60分钟，中间不要断电

刷完之后，在Jetson上验证：

sudo jetson_release
python3 -c "import torch; print(torch.cuda.is_available())"

4.4 RKNN Toolkit安装

如果你用的是瑞芯微的NPU（比如RK3588、RK3568），那就要用RKNN Toolkit。这个工具负责把模型转换成NPU能识别的格式。

安装步骤：

# 1. 安装依赖
sudo apt-get install libxslt1-dev libxml2-dev

# 2. 下载RKNN Toolkit
git clone https://github.com/rockchip-linux/rknn-toolkit2.git

# 3. 安装Python包
cd rknn-toolkit2/rknn-toolkit2/packages
pip install rknn_toolkit2-1.6.0-cp310-cp310-linux_x86_64.whl

验证：

python -c "from rknn.api import RKNN; print('RKNN Toolkit installed')"

4.5 整体工具链关系图

下面这张图，把整个工具链的流转关系画清楚了。从训练到部署，每一步用什么工具，一目了然。

4.6 环境验证清单

所有工具装完之后，我建议你跑一遍这个验证脚本。确保每个环节都没问题，再往下走。

#!/bin/bash
echo "=== 环境验证 ==="

echo "1. PyTorch"
python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA {torch.cuda.is_available()}')"

echo "2. ONNX"
python -c "import onnx; print(f'ONNX {onnx.__version__}')"

echo "3. TensorRT"
python -c "import tensorrt as trt; print(f'TensorRT {trt.__version__}')"

echo "4. RKNN"
python -c "from rknn.api import RKNN; print('RKNN OK')"

echo "=== 验证完成 ==="