2、环境配置：CUDA、cuDNN、PyTorch安装，Docker镜像构建，常见环境问题排查

环境配置这件事，说难不难，说简单也真能卡你三天。我见过太多同学模型写好了，结果一跑就报错，最后发现是CUDA版本和PyTorch没对齐。咱们这一章就把这事彻底理清楚。

2.1 CUDA与cuDNN：GPU加速的基石

CUDA是NVIDIA的并行计算平台，cuDNN是专门为深度学习优化的加速库。说白了，CUDA让GPU能干活，cuDNN让干活效率翻倍。

版本选择策略：

PyTorch版本	推荐CUDA版本	推荐cuDNN版本	备注
1.10 - 1.13	11.3 / 11.6	8.2 / 8.4	YOLOv5常见搭配
2.0 - 2.1	11.7 / 11.8	8.5 / 8.7	YOLOv8推荐
2.2+	12.1	8.9	最新特性支持

我个人习惯是：先查PyTorch官网的安装命令，看它默认捆绑哪个CUDA版本，然后照着装。这样最稳。

安装步骤（Linux为例）

1. 检查驱动：nvidia-smi 确认GPU被识别，驱动版本≥450
2. 下载CUDA Toolkit：从NVIDIA官网选对应版本，runfile方式安装
3. 配置环境变量：在 ~/.bashrc 添加 export PATH=/usr/local/cuda/bin:$PATH
4. 验证：nvcc --version 显示版本号即成功
5. 安装cuDNN：下载对应版本的tar包，解压后复制到CUDA目录

# cuDNN安装示例
tar -xzvf cudnn-linux-x86_64-8.7.0.84_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/
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

2.2 PyTorch安装：选对渠道很重要

PyTorch安装其实就一条命令的事，但坑都在细节里。

推荐安装方式：

• pip安装：最常用，PyTorch官方推荐
• conda安装：适合管理多个环境，但包可能不是最新
• 源码编译：除非你有特殊需求，否则别碰，编译一次够你喝一壶

# CUDA 11.8 + PyTorch 2.0
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 验证安装
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

如果输出 True，恭喜你，环境通了。如果是 False，别慌，往下看排查部分。

2.3 Docker镜像构建：一次配置，到处运行

Docker是解决环境一致性的终极方案。我团队现在所有训练任务都跑在Docker里，再也不用担心"在我机器上能跑啊"这种问题了。

Dockerfile示例：

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

# 安装基础依赖
RUN apt-get update && apt-get install -y \
    python3.9 python3-pip git wget \
    && rm -rf /var/lib/apt/lists/*

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

# 安装YOLO依赖
RUN pip3 install ultralytics opencv-python matplotlib

# 设置工作目录
WORKDIR /workspace

构建与运行：

# 构建镜像
docker build -t yolo-env:latest .

# 运行容器（挂载数据和代码）
docker run --gpus all -it \
    -v /path/to/data:/workspace/data \
    -v /path/to/code:/workspace/code \
    yolo-env:latest

2.4 常见环境问题排查

这部分是我最想分享的。踩过的坑多了，自然就知道怎么避。

问题1：torch.cuda.is_available() 返回 False

• 检查驱动：nvidia-smi 是否正常输出
• 检查CUDA版本：PyTorch编译时的CUDA版本和系统安装的是否匹配
• 检查权限：ls -l /dev/nvidia* 确认设备文件可访问

问题2：训练时显存不足（OOM）

• 降低batch size（最直接的方法）
• 开启梯度累积：accumulate = 4 等效于batch size扩大4倍
• 使用混合精度训练：model.half() 或AMP

问题3：cuDNN版本不匹配

• 错误特征：undefined symbol 或 libcudnn.so.x: cannot open shared object file
• 解决方案：用 ldd /path/to/torch/lib/libtorch_cuda.so | grep cudnn 查看实际链接的cuDNN版本

2.5 环境配置知识体系

下面这张图帮你理清整个环境配置的脉络：

2.6 快速验证脚本

每次配完环境，我建议你跑一遍这个脚本，全绿才算完事：

import torch
import sys

print(f"Python版本: {sys.version}")
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA是否可用: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"CUDA版本: {torch.version.cuda}")
    print(f"cuDNN版本: {torch.backends.cudnn.version()}")
    print(f"GPU数量: {torch.cuda.device_count()}")
    print(f"当前GPU: {torch.cuda.get_device_name(0)}")
    
    # 实际跑一下
    x = torch.randn(100, 3, 640, 640).cuda()
    print(f"张量测试通过: {x.shape}")
else:
    print("警告: CUDA不可用，将使用CPU训练（不推荐）")