2. 深度学习环境搭建

说实话，环境搭建这事儿，看着简单，但坑是真不少。我见过太多同学模型写好了，结果跑不起来，最后发现是CUDA版本和PyTorch对不上。这一章，咱们就把这事彻底理清楚。

2.1 环境概览：你到底需要什么？

先看一张整体流程图，搞清楚每个组件的位置：

深度学习环境分层结构 YOLO 训练 / 推理脚本 PyTorch (深度学习框架) cuDNN (深度神经网络加速库) NVIDIA 显卡驱动 + CUDA Toolkit NVIDIA GPU (硬件) 核心原则 版本必须严格对齐 从上到下逐层依赖 先装驱动，再装框架 常见坑点 PyTorch 版本 ≠ CUDA 版本 驱动版本决定 CUDA 上限 Docker 可隔离环境冲突

每一层都依赖下一层。你想想看，如果驱动没装好，CUDA Toolkit装得再漂亮也没用。PyTorch调用cuDNN时，底层一报错，整个训练就崩了。

2.2 NVIDIA驱动安装

这是第一步，也是最容易忽略的一步。我个人习惯是，拿到新机器第一件事就是查驱动版本。

检查当前驱动：

nvidia-smi

如果显示GPU信息，说明驱动已装。如果报错，去NVIDIA官网下载对应驱动。

这里有个关键点：驱动版本决定了你能装多高的CUDA。比如驱动版本是470.x，那CUDA最高只能到11.4。我曾在项目中遇到过，客户服务器驱动太老，死活装不上CUDA 11.8，最后只能降级PyTorch版本。

驱动版本 支持的最高CUDA 推荐场景
470.x CUDA 11.4 老旧服务器、T4显卡
510.x CUDA 11.6 一般工作站
525.x CUDA 12.0 RTX 30/40系列
535.x CUDA 12.2 最新显卡、H100

注意：不要装最新的驱动！除非你确定你的PyTorch支持。我建议装「推荐/认证」版本，稳定第一。

2.3 CUDA Toolkit安装

驱动搞定后，装CUDA Toolkit。说白了，它就是给开发者用的工具包，包含了编译器、库文件等。

安装步骤：

  1. 去NVIDIA官网下载CUDA Toolkit（我推荐11.8，兼容性最好）
  2. 选择对应操作系统和架构
  3. 用runfile或deb包安装

我的习惯：用runfile安装，可以自定义安装路径。这样一台机器可以装多个CUDA版本，通过修改环境变量切换。

验证安装：

nvcc --version

如果显示版本号，说明CUDA Toolkit装好了。

2.4 cuDNN安装

cuDNN是NVIDIA专门为深度学习优化的加速库。PyTorch底层大量调用它来做卷积、池化等操作。

安装其实很简单：

  1. 去NVIDIA官网下载cuDNN（需要注册账号）
  2. 选择与CUDA版本匹配的cuDNN
  3. 解压后把文件复制到CUDA安装目录
# 示例：将cuDNN文件复制到CUDA目录
sudo cp cudnn-linux-x86_64-8.9.0.131_cuda11-archive/include/cudnn*.h /usr/local/cuda/include/
sudo cp cudnn-linux-x86_64-8.9.0.131_cuda11-archive/lib/libcudnn* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

我曾经踩过的坑：cuDNN版本和CUDA版本不匹配，训练时直接报「cuDNN error」。后来我养成了习惯，每次装完都跑一下cuDNN的sample程序验证。

2.5 PyTorch安装与验证

终于到框架层了。PyTorch的安装其实是最简单的，但版本选择有讲究。

安装命令（以CUDA 11.8为例）：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

为什么推荐用官网命令？因为PyTorch的wheel包已经编译好了cuDNN，你不需要手动再装一遍。但前提是，你的CUDA版本必须和PyTorch要求的匹配。

PyTorch版本 对应CUDA 稳定程度
1.13.x CUDA 11.7 非常稳定，推荐
2.0.x CUDA 11.8 稳定，新特性多
2.1.x CUDA 12.1 较新，部分库兼容性待验证

验证PyTorch是否能用GPU：

import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))

如果输出True和你的显卡型号，恭喜你，环境搭好了。

快速验证脚本：

import torch
x = torch.randn(3, 3).cuda()
print(x)
print("GPU is working!")

2.6 Docker环境配置

说实话，前面那些步骤，每换一台机器就要重来一遍，烦不烦？我反正是烦了。所以后来我全面转向Docker。

Docker的好处：

  • 环境隔离，一个容器一个项目
  • 一键部署，不用重复装驱动
  • 版本管理，想换CUDA就换镜像

安装Docker：

sudo apt update
sudo apt install docker.io
sudo systemctl start docker
sudo systemctl enable docker

安装NVIDIA Container Toolkit：

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 update
sudo apt install nvidia-container-toolkit
sudo systemctl restart docker

拉取PyTorch镜像并运行：

docker pull pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
docker run --gpus all -it --rm pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime python -c "import torch; print(torch.cuda.is_available())"

如果输出True，说明Docker环境也搞定了。

我的Dockerfile模板：

FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
RUN pip install opencv-python matplotlib tqdm
WORKDIR /workspace
COPY . .
CMD ["python", "train.py"]

这样每次换机器，只需要 docker build -t my-yolo . 然后 docker run --gpus all -it my-yolo，环境一模一样。

2.7 常见问题与避坑指南

  • nvidia-smi报错：驱动没装好，或者内核模块没加载。重启试试。
  • torch.cuda.is_available()返回False：PyTorch版本和CUDA版本不匹配。我建议用官网的pip命令重装。
  • Docker里找不到GPU：没装nvidia-container-toolkit，或者没加 --gpus all 参数。
  • cuDNN报错：版本不匹配，或者没复制到正确目录。用 ldconfig 刷新库缓存。

我曾经犯过的错：在一台机器上装了CUDA 11.8，但PyTorch是2.1版本默认支持CUDA 12.1，结果训练时各种诡异报错。后来我统一用PyTorch官网推荐的CUDA版本，再也没出过问题。

环境搭建这事，说白了就是「版本对齐」。驱动、CUDA、cuDNN、PyTorch，四者必须匹配。用Docker可以省去很多重复劳动，我强烈推荐。

← 上一章 📖 返回目录 下一章 →