第2章环境搭建与工具链

2、环境搭建与工具链：CUDA/cuDNN安装、TensorRT安装、ONNX Runtime配置、Docker环境

说实话，环境搭建这事儿，看着简单，但坑是真不少。我见过太多同学模型写好了，结果卡在环境配置上，一卡就是好几天。咱们这节课，我就把这几年来踩过的坑、总结的经验，一次性给你讲透。

2.1 CUDA与cuDNN：GPU加速的基石

CUDA是NVIDIA的并行计算平台。说白了，它就是让GPU干活的“翻译官”。cuDNN则是专门为深度学习优化的加速库。两者缺一不可。

版本匹配是头等大事

我曾经因为CUDA和cuDNN版本不匹配，整整排查了两天。后来发现，就是一个小版本号没对上。记住：CUDA版本必须与cuDNN版本严格对应。

安装步骤其实不复杂，我习惯这么做：

先确认你的GPU型号——运行 nvidia-smi 查看驱动版本
去NVIDIA官网下载CUDA Toolkit——我推荐用runfile方式安装，灵活可控
配置环境变量——把CUDA路径加到 PATH 和 LD_LIBRARY_PATH 里
安装cuDNN——解压后把文件复制到CUDA目录下

# 我的标准安装流程
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
sudo sh cuda_11.8.0_520.61.05_linux.run

# 配置环境变量（加到 ~/.bashrc）
export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH

# cuDNN安装
tar -xzvf cudnn-linux-x86_64-8.9.0.131_cuda11-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include/
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64/

小技巧

安装完后，运行 nvcc --version 和 cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR 验证版本。我每次装完都会跑一遍，省得后面出问题。

2.2 TensorRT：推理加速的利器

TensorRT是NVIDIA的推理优化引擎。它能做什么？简单说，就是把你的模型“压缩”成最适合GPU运行的格式。我做过测试，YOLOv5用TensorRT加速后，推理速度能提升3-5倍。

安装TensorRT，我个人推荐用tar包方式。为什么？因为deb包容易和系统库冲突，tar包干净利落。

# 下载TensorRT（需要NVIDIA账号）
wget https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/8.6.1/tars/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz

# 解压
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz

# 配置环境变量
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/TensorRT-8.6.1.6/lib

# 安装Python包
cd TensorRT-8.6.1.6/python
pip install tensorrt-8.6.1-cp38-none-linux_x86_64.whl

关键点

TensorRT版本必须和CUDA版本匹配。比如CUDA 11.8对应TensorRT 8.6.x。我见过有人用CUDA 12.0配TensorRT 8.4，结果各种报错。嗯，这里要注意。

2.3 ONNX Runtime：跨平台的推理引擎

ONNX Runtime是个好东西。它支持多种硬件后端，CPU、GPU、甚至手机端都能跑。我一般用它来做模型验证和快速原型开发。

配置起来很简单：

# CPU版本
pip install onnxruntime

# GPU版本（需要CUDA和cuDNN）
pip install onnxruntime-gpu

# 验证安装
python -c "import onnxruntime; print(onnxruntime.get_device())"

你想想看，ONNX Runtime最大的优势是什么？是它的算子兼容性。很多模型转TensorRT会报错，但ONNX Runtime基本都能跑。我有个项目，客户要求同时支持NVIDIA和AMD显卡，ONNX Runtime就是最佳选择。

实战建议

先用ONNX Runtime跑通模型，再转TensorRT做优化。这样能快速定位问题，避免在优化阶段浪费时间。

2.4 Docker环境：一次配置，到处运行

Docker是我最推荐的部署方式。为什么？因为它解决了“在我电脑上能跑”这个千古难题。

我习惯用NVIDIA官方提供的Docker镜像作为基础：

# 拉取官方镜像
docker pull nvcr.io/nvidia/tensorrt:23.12-py3

# 运行容器
docker run --gpus all -it --rm \
  -v /path/to/your/project:/workspace \
  nvcr.io/nvidia/tensorrt:23.12-py3

# 在容器内安装ONNX Runtime
pip install onnxruntime-gpu

注意

Docker容器默认没有图形界面。如果你需要可视化调试，记得加上 --net=host --env="DISPLAY" --volume="$HOME/.Xauthority:/root/.Xauthority:rw" 参数。

Docker的好处太多了。我举个例子：你开发时用CUDA 11.8，但客户环境是CUDA 12.0。怎么办？用Docker打包你的环境，客户直接拉镜像就能跑。省心省力。

2.5 环境验证与常见问题

装完环境，一定要做验证。我有个检查清单：

检查项	命令	预期结果
CUDA版本	`nvcc --version`	显示11.8或对应版本
cuDNN版本	`cat /usr/local/cuda/include/cudnn_version.h \| grep CUDNN_MAJOR`	显示8.9或对应版本
TensorRT版本	`python -c "import tensorrt; print(tensorrt.__version__)"`	显示8.6或对应版本
ONNX Runtime	`python -c "import onnxruntime; print(onnxruntime.get_device())"`	显示GPU

避坑指南

我曾经遇到一个问题：TensorRT能导入，但推理时一直报错。后来发现是cuDNN的so文件没复制全。记住：cuDNN的lib文件必须全部复制到CUDA的lib64目录下，一个都不能少。

好了，环境搭建这部分就讲到这里。下一章咱们开始实战——用TensorRT优化YOLO模型。到时候你会发现，今天花时间把环境配好，后面能省下大把时间。