📚 课程目录第2章环境搭建

2、环境搭建：CUDA与cuDNN安装、PyTorch环境配置、ONNX与ONNX Runtime安装、TensorRT安装、Docker环境搭建

说实话，环境搭建是很多同学入门YOLO部署的第一道坎。我见过太多人在这一步卡了两三天，最后发现是CUDA版本和PyTorch没对齐。这一章，咱们就把这些坑一个个填平。

我的习惯是：先搭好底层驱动，再往上装框架，最后用Docker把环境固化。这样不管换机器还是重装系统，都能快速恢复。

核心原则：版本对齐

CUDA、cuDNN、PyTorch、TensorRT这四个东西，版本必须严格匹配。差一个小版本，可能就报错给你看。

2.1 CUDA与cuDNN安装

CUDA是NVIDIA的并行计算平台，cuDNN是它的深度学习加速库。说白了，CUDA是发动机，cuDNN是涡轮增压器。

2.1.1 检查硬件兼容性

先确认你的显卡支持什么版本的CUDA。用这个命令：

nvidia-smi

输出里会显示Driver Version和CUDA Version。注意，这个CUDA Version是驱动支持的最高版本，不是你已经装了的版本。

显卡架构	计算能力	推荐CUDA版本
Ampere (RTX 30系列)	8.0+	CUDA 11.x / 12.x
Turing (RTX 20系列)	7.5	CUDA 10.x / 11.x
Volta (V100)	7.0	CUDA 10.x / 11.x
Maxwell/Pascal (GTX 10系列)	5.x / 6.x	CUDA 10.x

2.1.2 安装CUDA

我推荐用runfile方式安装，这样能自己选择安装路径，方便管理多个版本。

# 下载CUDA 11.8（以这个版本为例）
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

安装时注意：不要勾选Driver，除非你的驱动版本太旧。我只装CUDA Toolkit本身。

我的经验：装完CUDA后，一定要配环境变量。我习惯把多个CUDA版本放在/usr/local/cuda-11.8、/usr/local/cuda-12.1这样的目录下，然后用软链接切换。

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

# 验证安装
nvcc --version

2.1.3 安装cuDNN

cuDNN需要去NVIDIA官网下载，要注册账号。下载后解压，把文件复制到CUDA目录就行。

# 解压cuDNN
tar -xzvf cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz

# 复制到CUDA目录
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda-11.8/include/
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda-11.8/lib64/

# 设置权限
sudo chmod a+r /usr/local/cuda-11.8/include/cudnn*.h
sudo chmod a+r /usr/local/cuda-11.8/lib64/libcudnn*

# 验证
cat /usr/local/cuda-11.8/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

注意：cuDNN版本必须和CUDA版本对应。比如CUDA 11.8就用cuDNN 8.9.x for CUDA 11.x。我曾经因为cuDNN版本不对，训练时直接报"undefined symbol"错误，排查了半天。

2.2 PyTorch环境配置

PyTorch的安装其实是最简单的，但也是最容易出错的——因为版本匹配问题。

2.2.1 创建虚拟环境

我强烈建议用conda管理环境。每个项目一个环境，互不干扰。

# 创建Python 3.10环境
conda create -n yolo_deploy python=3.10
conda activate yolo_deploy

2.2.2 安装PyTorch

去PyTorch官网用版本匹配器生成安装命令。比如CUDA 11.8对应：

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，说明CUDA可用。如果False，检查一下是不是装了CPU版本的PyTorch。

避坑指南：我曾经用pip install torch直接装，结果装成了CPU版。后来养成习惯，每次都指定--index-url参数。

2.3 ONNX与ONNX Runtime安装

ONNX是模型交换格式，ONNX Runtime是推理引擎。YOLO模型转ONNX后，可以用ONNX Runtime做快速验证。

2.3.1 安装ONNX

pip install onnx==1.15.0

2.3.2 安装ONNX Runtime

ONNX Runtime分CPU版和GPU版。做YOLO部署，肯定用GPU版。

# GPU版（CUDA 11.x）
pip install onnxruntime-gpu==1.17.0

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

输出应该是"GPU"。如果是"CPU"，检查一下CUDA环境变量有没有配好。

组件	推荐版本	说明
ONNX	1.15.x	用于模型导出和检查
ONNX Runtime GPU	1.17.x	GPU推理，依赖CUDA
ONNX Runtime CPU	1.17.x	CPU推理，调试用

2.4 TensorRT安装

TensorRT是NVIDIA的推理优化引擎。YOLO模型转TensorRT后，推理速度能提升2-5倍。

2.4.1 下载TensorRT

去NVIDIA官网下载TensorRT。注意选对CUDA版本和系统架构。

# 下载TensorRT 8.6.1 for CUDA 11.8
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/8.6.1/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
sudo mv TensorRT-8.6.1.6 /usr/local/

2.4.2 配置环境变量

# 加到~/.bashrc
export LD_LIBRARY_PATH=/usr/local/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH
export PATH=/usr/local/TensorRT-8.6.1.6/bin:$PATH

2.4.3 安装Python包

# 进入TensorRT的Python目录
cd /usr/local/TensorRT-8.6.1.6/python

# 安装对应Python版本的wheel
pip install tensorrt-8.6.1-cp310-none-linux_x86_64.whl

# 验证
python -c "import tensorrt; print(tensorrt.__version__)"

注意：TensorRT的Python包版本必须和你的Python版本匹配。cp310就是Python 3.10。我见过有人装了cp39的包，Python是3.10，结果import就报错。

2.5 Docker环境搭建

Docker是环境管理的终极方案。一次配置，到处运行。

2.5.1 安装Docker

# Ubuntu安装Docker
sudo apt-get update
sudo apt-get install docker.io

# 启动Docker服务
sudo systemctl start docker
sudo systemctl enable docker

# 验证
docker --version

2.5.2 安装NVIDIA Container Toolkit

这个工具让Docker容器能访问GPU。

# 添加NVIDIA仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

# 安装
sudo apt-get update
sudo apt-get install -y nvidia-docker2

# 重启Docker
sudo systemctl restart docker

# 验证GPU是否可用
docker run --rm --gpus all nvidia/cuda:11.8.0-base nvidia-smi

2.5.3 构建YOLO部署镜像

我习惯用Dockerfile把环境固化下来。

# Dockerfile
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04

# 安装Python和pip
RUN apt-get update && apt-get install -y python3 python3-pip

# 安装PyTorch
RUN pip3 install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118

# 安装ONNX和ONNX Runtime
RUN pip3 install onnx==1.15.0 onnxruntime-gpu==1.17.0

# 安装TensorRT
COPY TensorRT-8.6.1.6 /opt/TensorRT-8.6.1.6
ENV LD_LIBRARY_PATH=/opt/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH
RUN pip3 install /opt/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp38-none-linux_x86_64.whl

# 设置工作目录
WORKDIR /workspace

构建镜像：

docker build -t yolo-deploy:1.0 .

我的习惯：Docker镜像打好后，我会push到私有仓库。这样团队里其他人直接pull就能用，不用重复装环境。有一次服务器重装系统，我花了10分钟拉镜像，环境就全回来了。

2.6 环境验证与知识图谱

所有组件装完后，跑一个完整的验证脚本：

python -c "
import torch
import onnx
import onnxruntime
import tensorrt

print(f'PyTorch: {torch.__version__}')
print(f'CUDA可用: {torch.cuda.is_available()}')
print(f'ONNX: {onnx.__version__}')
print(f'ONNX Runtime: {onnxruntime.__version__}')
print(f'TensorRT: {tensorrt.__version__}')
"

下面这张图，帮你理清整个环境搭建的知识体系：

环境搭好了，后面的事情就顺了。记住：版本对齐是环境搭建的第一要务。我每次搭环境都会列一个版本清单，逐个核对，从不跳过这一步。