使用 Docker 在 PyTorch 环境中训练模型

在机器学习和深度学习任务中，使用 Docker 可以方便地构建和管理环境，特别是在涉及复杂依赖关系和 GPU 加速时。本文介绍如何基于官方 PyTorch 镜像构建训练环境，并在容器中运行训练脚本。

准备工作

首先编写一个 Dockerfile，描述镜像应该包含的依赖和启动命令。官方 PyTorch 镜像已经内置对应版本的 torch、torchvision、CUDA 和 cuDNN，通常不要在镜像里再次 pip install torch torchvision，否则容易把 CUDA 版本和 PyTorch 版本装乱。

下面示例使用一个明确版本的官方镜像。实际项目中应根据服务器 NVIDIA 驱动版本选择匹配的 CUDA 镜像标签。

# 使用官方 PyTorch 镜像作为基础镜像
FROM pytorch/pytorch:2.13.0-cuda12.6-cudnn9-devel

# 设置工作目录
WORKDIR /app

# 先复制依赖声明，便于利用 Docker 构建缓存
COPY requirements.txt /app/requirements.txt

# 只安装业务额外依赖，不重复安装 torch / torchvision
RUN pip install --no-cache-dir -r requirements.txt

# 复制训练脚本
COPY train.py /app/train.py

# 启动应用程序
CMD ["python", "train.py"]

requirements.txt 中只放业务依赖，例如：

pandas
tqdm
seaborn
requests

其中，train.py 是训练脚本，也放在同一目录。

构建 Docker 镜像

在 Dockerfile 所在目录下，打开终端并运行以下命令来构建 Docker 镜像：

1	docker build -t test_train .

test_train 是生成的 Docker 镜像名称。

运行 Docker 容器

构建完成后，我们可以使用以下命令来运行 Docker 容器，并在其中执行训练脚本：

1	docker run --gpus all -it --rm --shm-size=4g test_train

在这个命令中，--gpus all 用于启用 GPU 支持，-it 表示以交互模式运行容器，--rm 表示容器停止后立即删除，--shm-size 用于设置容器共享内存大小。训练脚本如果使用 DataLoader(num_workers>0) 或较大的 batch，过小的共享内存可能导致 worker 异常退出。