环境准备：Docker基础、容器化训练环境、GPU驱动与CUDA配置

说实话，很多做算法的人一听到「环境配置」就头疼。我完全理解——我自己早期也在这上面栽过跟头。有一次花了两天配好的环境，换台机器全崩了。后来我学乖了，老老实实用 Docker 把训练环境打包好，走到哪都能跑。

这一章，我们就来搞定这件事。说白了，就是三块内容：Docker 怎么用、怎么把训练环境装进容器、以及 GPU 驱动和 CUDA 怎么配才能让容器真正用上显卡。

为什么非要用 Docker？

你想想看，一个典型的深度学习项目，依赖的东西有多少？Python 版本、CUDA 版本、cuDNN、PyTorch 或 TensorFlow、各种 pip 包……每个人机器不一样，装出来的效果也不一样。这就是所谓的「在我机器上能跑」的尴尬。

Docker 解决的就是这个问题。它把你的整个环境打包成一个镜像，别人拉下来就能用，一模一样。我在团队里推 Docker 的时候，运维同事看我的眼神都不一样了——终于不用半夜爬起来帮人调环境了。

核心概念速览：

  • 镜像（Image）：一个只读的模板，包含操作系统、软件、代码等所有东西
  • 容器（Container）：镜像的运行实例，可以启动、停止、删除
  • Dockerfile：构建镜像的配方文件，告诉 Docker 怎么一步步搭环境
  • 仓库（Registry）：存放镜像的地方，比如 Docker Hub

Docker 基础操作

先别急着装东西，我们看看最常用的几个命令。我平时 80% 的时间就靠这几个命令活着。

# 拉取镜像
docker pull nvidia/cuda:12.1.0-devel-ubuntu22.04

# 查看本地镜像
docker images

# 运行容器（交互模式）
docker run -it --name my_train_env nvidia/cuda:12.1.0-devel-ubuntu22.04 /bin/bash

# 查看运行中的容器
docker ps

# 停止容器
docker stop my_train_env

# 启动已停止的容器
docker start -ai my_train_env

# 删除容器
docker rm my_train_env

# 删除镜像
docker rmi nvidia/cuda:12.1.0-devel-ubuntu22.04

嗯，这里有个小细节：-it 是交互模式，-d 是后台运行。我刚开始经常搞混，后来记住一个口诀——「要进去干活用 -it，要跑服务用 -d」。

编写 Dockerfile：把训练环境装进镜像

手动装环境太累了，而且容易出错。更好的做法是写一个 Dockerfile，把每一步都写清楚。以后改环境，改 Dockerfile 重新构建就行。

下面是我常用的一个训练环境 Dockerfile，你直接拿去改改就能用。

# 基础镜像：选一个带 CUDA 的 Ubuntu
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04

# 设置非交互模式，避免安装时卡住
ENV DEBIAN_FRONTEND=noninteractive

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    git \
    wget \
    vim \
    && rm -rf /var/lib/apt/lists/*

# 设置 Python 软链接
RUN ln -s /usr/bin/python3.10 /usr/bin/python

# 升级 pip
RUN pip3 install --upgrade pip

# 安装深度学习框架
RUN pip3 install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121

# 安装常用库
RUN pip3 install \
    numpy \
    pandas \
    scikit-learn \
    matplotlib \
    tensorboard \
    jupyter

# 创建工作目录
WORKDIR /workspace

# 默认命令
CMD ["/bin/bash"]

我的经验：基础镜像选 devel 版本而不是 runtime 版本。虽然 devel 更大，但它包含了编译工具链。我遇到过好几次 runtime 镜像里没有 gcc，装某些包时直接报错。省那几百兆空间，不值得。

构建镜像的命令也很简单：

docker build -t my_training_env:v1 .

等几分钟，镜像就构建好了。以后换机器，只要把这个 Dockerfile 拿过去重新 build，或者把镜像 push 到仓库再 pull，环境就一模一样。

GPU 驱动与 CUDA 配置

这是最容易出问题的地方。我见过太多人卡在这一步——容器跑起来了，但 nvidia-smi 报错，或者 PyTorch 检测不到 GPU。

宿主机先装好驱动

首先，你的宿主机（就是物理机或云服务器）上必须装好 NVIDIA 驱动。这个驱动是基础，容器里的 CUDA 依赖它来跟显卡通信。

# 检查驱动是否安装
nvidia-smi

# 如果没装，用这个命令装（Ubuntu 为例）
sudo apt-get install nvidia-driver-535

# 装完重启
sudo reboot

驱动装好后，nvidia-smi 应该能正常显示 GPU 信息。注意看右上角的 CUDA Version，那是驱动支持的 CUDA 版本。比如显示 CUDA Version: 12.1，说明这个驱动可以配合 CUDA 12.1 及以下版本使用。

我曾经踩过的坑：驱动版本太老，容器里装的是 CUDA 12.1，但驱动只支持到 11.8。结果容器启动后，PyTorch 报错说找不到 CUDA 驱动。后来我学乖了——先看驱动支持的 CUDA 版本，再选对应的基础镜像。

安装 NVIDIA Container Toolkit

光有驱动还不够，Docker 容器默认是看不到 GPU 的。你需要装一个工具，让容器能访问宿主的 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-container-toolkit

# 重启 Docker
sudo systemctl restart docker

装完之后，启动容器时加上 --gpus all 参数，容器就能用 GPU 了。

docker run -it --gpus all my_training_env:v1 nvidia-smi

如果能看到 GPU 信息，说明配置成功了。

容器化训练环境的完整流程

把上面这些串起来，就是一个完整的流程。我画了一张图，帮你理清思路。

容器化训练环境搭建流程 宿主机准备 安装 NVIDIA 驱动 安装 Docker 安装 nvidia-container-toolkit 编写 Dockerfile docker build 构建镜像 docker run --gpus all nvidia-smi 验证 💡 验证通过后，训练环境就准备好了 后续所有训练都在容器内进行

常用镜像对比

选基础镜像的时候，很多人会纠结。我整理了一个表格，帮你快速决策。

镜像类型 大小 包含内容 适用场景
nvidia/cuda:12.1.0-base-ubuntu22.04 ~1.2GB CUDA 运行时库 只需要跑已编译好的程序
nvidia/cuda:12.1.0-runtime-ubuntu22.04 ~1.8GB CUDA 运行时 + 部分工具 大多数训练场景
nvidia/cuda:12.1.0-devel-ubuntu22.04 ~3.5GB 完整 CUDA 工具链 + 编译工具 需要编译自定义 CUDA 算子
pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel ~5.0GB PyTorch + CUDA + cuDNN 预装好 快速上手，省去自己装 PyTorch

我的建议：刚开始用 pytorch/pytorch 官方镜像最省事。等你对环境熟悉了，再自己从 nvidia/cuda:devel 开始定制。我自己现在都是基于 devel 镜像搭，因为经常要编译一些第三方 CUDA 扩展。

验证环境是否正常

环境配好之后，一定要验证。我一般跑下面这段代码：

import torch

# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")

# 检查 GPU 数量
print(f"GPU count: {torch.cuda.device_count()}")

# 检查当前 GPU 名称
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")
    
# 跑一个简单的张量运算
x = torch.randn(3, 3).cuda()
y = torch.randn(3, 3).cuda()
z = x + y
print(f"Tensor operation on GPU: {z}")

如果输出显示 CUDA available: True，并且能看到 GPU 名称，那就说明环境配好了。

注意：如果 torch.cuda.is_available() 返回 False，别慌。先检查容器启动时有没有加 --gpus all，再检查宿主机驱动版本和容器内 CUDA 版本是否兼容。我遇到过最坑的一次是 Docker 版本太老，不支持 --gpus 参数，升级 Docker 就好了。

日常开发小技巧

最后分享几个我平时用的小技巧，能省不少时间。

  • 挂载代码目录：用 -v /host/path:/workspace 把宿主机代码目录挂载到容器里，这样改代码不用重新构建镜像
  • 端口映射：用 -p 8888:8888 把容器的 Jupyter 端口映射出来，方便在浏览器里用
  • 容器命名：用 --name 给容器起个有意义的名字，不然一堆随机字符串根本分不清
  • 自动删除：测试时加 --rm，容器退出后自动删除，不会留下一堆废弃容器
# 一个完整的启动命令示例
docker run -it --gpus all \
  --name my_training \
  -v /home/user/project:/workspace \
  -p 8888:8888 \
  --rm \
  my_training_env:v1

好了，环境准备这块就讲完了。你按照这个流程走一遍，应该不会出大问题。如果卡住了，多半是驱动版本或者 Docker 权限的问题——这两样我踩过的坑最多，你留意一下就行。

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