Docker 与 GPU 训练

以下是 Docker 与 GPU 训练 相关的核心命令和操作指南，涵盖容器启动、资源监控和调试技巧：

1. 启用 GPU 支持的 Docker

安装 NVIDIA 容器工具包（必须步骤）

bash 复制代码

# 添加源并安装
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 -y nvidia-docker2
sudo systemctl restart docker

验证安装

bash 复制代码

docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

若看到与宿主机相同的 GPU 信息输出，则配置成功。

2. 启动带 GPU 的容器

基础命令

bash 复制代码

# 使用所有GPU
docker run --gpus all -it nvidia/cuda:11.0-base bash

# 指定GPU数量（如2卡）
docker run --gpus 2 -it nvidia/cuda:11.0-base bash

# 指定具体GPU设备（如仅用第0,1号卡）
docker run --gpus '"device=0,1"' -it nvidia/cuda:11.0-base bash

完整训练示例（PyTorch）

bash 复制代码

docker run -d --name pytorch_train \
  --gpus all \
  -v $(pwd)/data:/data \
  -v $(pwd)/code:/code \
  pytorch/pytorch:1.9.0-cuda11.1-cudnn8-runtime \
  python /code/train.py

3. 监控 GPU/CPU 资源

查看容器内 GPU 使用

bash 复制代码

# 进入容器后直接运行
nvidia-smi

# 或宿主机查看所有容器GPU占用
docker stats --format "table {{.Name}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.NetIO}}"

实时监控工具

bash 复制代码

# 安装nvtop（推荐）
sudo apt install nvtop
nvtop  # 类似htop的GPU监控界面

查看容器绑定的 GPU 设备

bash 复制代码

docker inspect <容器ID> | grep -i nvidia

4. 调试与常见问题

问题1：容器内找不到 GPU

现象：nvidia-smi 报错或无输出。

解决：

bash 复制代码

# 检查Docker是否支持GPU
docker info | grep -i runtime
# 应显示：runtimes: nvidia, runc

# 重新加载驱动
sudo modprobe nvidia

问题2：GPU 显存不足

优化方案 ：

bash 复制代码

# 启动时限制显存（如每卡最多8G）
docker run --gpus all --cpuset-cpus 0-3 --memory 16g -it ...

问题3：多容器共享 GPU

使用 MIG 技术 （仅限A100/H100等显卡）：

bash 复制代码

# 将GPU划分为多个实例
nvidia-smi mig -cgi 1g.5gb -C
# 启动容器时指定MIG实例
docker run --gpus device=0:0 ...

5. 高级配置

指定 CUDA 版本

bash 复制代码

# 使用特定CUDA版本的镜像
docker run --gpus all nvidia/cuda:11.3.0-base nvcc --version

限制 CPU 核心

bash 复制代码

# 绑定到0-3号CPU核心
docker run --gpus all --cpuset-cpus 0-3 ...

容器内启用 GPU 计算验证

bash 复制代码

docker run --gpus all -it tensorflow/tensorflow:2.6.0-gpu \
  python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

6. 实战示例

启动 Jupyter Lab 并启用 GPU

bash 复制代码

docker run -d --name jupyter_gpu \
  --gpus all \
  -p 8888:8888 \
  -v $(pwd):/workspace \
  tensorflow/tensorflow:2.6.0-gpu-jupyter \
  jupyter lab --ip=0.0.0.0 --allow-root

访问 http://localhost:8888 并运行 !nvidia-smi 验证 GPU。

附：命令速查表

场景	命令
启用GPU支持	`docker run --gpus all ...`
指定GPU数量	`docker run --gpus 2 ...`
监控GPU	`nvidia-smi` / `nvtop`
调试驱动	`docker info
限制显存	`--gpus all --memory 16g`

掌握这些命令后，即可在 Docker 中高效运行 GPU 训练任务！