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

2. 训练环境搭建：深度学习框架选择、GPU驱动与CUDA安装、Docker环境配置

环境搭建这事儿，说简单也简单，说坑也多。我带过不少团队，发现很多人把时间花在调模型上，结果环境配不好，一跑就崩。今天咱们就把这块捋清楚。

2.1 深度学习框架：PyTorch vs TensorFlow

选框架就像选工具，顺手最重要。我个人习惯是看项目场景来决定。

对比维度	PyTorch	TensorFlow
调试体验	动态图，print就能看中间结果	静态图为主，调试相对麻烦
部署生态	TorchScript / ONNX 导出方便	TF Serving / TFLite 成熟
社区活跃度	学术界首选，代码多	工业界广泛，生产案例多
学习曲线	更接近Python原生，上手快	概念多，需要适应期

我自己的经验是：做研究、快速验证原型，PyTorch 更顺手；做大规模生产部署、需要完整MLOps链路，TensorFlow 的生态更全。但说实话，现在两者差距在缩小，PyTorch 2.0 的 torch.compile 也追上来了。

我的建议： 如果你刚入门，从 PyTorch 开始。等需要部署到移动端或嵌入式设备时，再学 TensorFlow Lite 也不迟。

2.2 GPU驱动与CUDA安装

这块是重灾区。我见过太多人卡在驱动版本不匹配上。说白了，就是你的显卡驱动、CUDA 版本、cuDNN 版本、框架版本，四者必须对齐。

2.2.1 确认硬件

先看看你手头的显卡。NVIDIA 的卡才能用 CUDA。用这个命令查：

nvidia-smi

输出里会显示驱动版本、CUDA 版本、显存信息。如果没输出，说明驱动没装好。

2.2.2 驱动安装

驱动版本不是越新越好。我踩过坑：装了最新驱动，结果 CUDA 11.8 不支持。后来学乖了，先确定 CUDA 版本，再装对应驱动。

推荐去 NVIDIA 官网下载，别用系统自带的驱动管理器。选的时候注意：

笔记本用户选 Notebook 版本
桌面用户选 Desktop 版本
Linux 用户选 runfile 或 deb 包

2.2.3 CUDA 与 cuDNN 安装

CUDA 安装其实不复杂，关键是版本匹配。我整理了一个常用对照表：

PyTorch 版本	推荐 CUDA	推荐 cuDNN
1.13.x	11.7	8.5
2.0.x	11.8	8.7
2.1.x	12.1	8.9

安装步骤：

下载 CUDA Toolkit（选 runfile 或 deb）
安装时不要装驱动（已经装过了）
配置环境变量 PATH 和 LD_LIBRARY_PATH
下载 cuDNN，解压后复制到 CUDA 目录

我曾经踩过的坑： 装完 CUDA 后没重启终端，环境变量没生效。跑 torch.cuda.is_available() 返回 False，排查了半天。记住：装完一定要 source ~/.bashrc 或者新开一个终端。

2.3 Docker环境配置

Docker 是环境管理的终极方案。你想想看，团队里每个人系统不一样，依赖版本不一样，用 Docker 就能保证「在我机器上能跑」变成「在哪都能跑」。

2.3.1 为什么用 Docker

隔离性： 不会污染宿主机环境
可复现： 一份 Dockerfile，到处运行
版本管理： 不同项目用不同 CUDA 版本

2.3.2 安装 Docker 和 nvidia-docker

Docker 安装很简单，但要让 Docker 用 GPU，需要装 nvidia-docker2。

# 安装 Docker
curl -fsSL https://get.docker.com -o get-docker.sh
sudo sh get-docker.sh

# 安装 nvidia-docker2
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
sudo systemctl restart docker

2.3.3 拉取 PyTorch 镜像

NVIDIA 官方提供了带 CUDA 和 cuDNN 的 PyTorch 镜像，省去自己配的麻烦。

# 拉取 PyTorch 2.1 + CUDA 12.1 镜像
docker pull pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel

# 运行容器，挂载代码目录
docker run --gpus all -it --rm \
  -v /home/user/project:/workspace \
  pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel

小技巧： 用 --gpus all 让容器使用所有 GPU。如果只想用某一块，改成 --gpus '"device=0"'。

2.4 环境验证

装完别急着跑模型，先验证环境是否正常。我一般跑这几行：

import torch
print("PyTorch 版本:", torch.__version__)
print("CUDA 可用:", torch.cuda.is_available())
print("GPU 数量:", torch.cuda.device_count())
print("当前 GPU:", torch.cuda.get_device_name(0))

# 跑一个简单测试
x = torch.randn(3, 3).cuda()
print("Tensor 在 GPU 上:", x.device)

如果输出正常，说明环境搭好了。如果 torch.cuda.is_available() 返回 False，别慌，按下面步骤排查：

检查 nvidia-smi 是否正常
检查 CUDA 版本是否匹配
检查 Docker 是否加了 --gpus all
检查环境变量 LD_LIBRARY_PATH

  核心要点： 环境搭建的关键是版本对齐。驱动、CUDA、cuDNN、框架，四者必须匹配。用 Docker 可以大幅降低环境问题带来的痛苦。

嗯，环境搭建这块就这些。记住一个原则：版本对齐。别贪新，别图省事，每一步验证好再往下走。这样后面训练和部署才能顺风顺水。