1. 镜像 Tag 标识的含义
base/cuda: 包括 CUDA 运行时
runtime: 在 base 的基础上，新增了 CUDA math 库和 NCCL、cuDNN 运行时
devel: 在 runtime 的基础上，新增了头文件和用于构建 CUDA 镜像的开发工具，对于多阶段构建特别有用
cuddn: 在上面基础上，新增了 cuDNN 神经网络加速库
py3: Python 3 环境
2. CUDA 镜像
| 镜像 | AMD64 镜像大小 | ARM64 镜像大小 |
|---|
| nvidia/cuda:12.3.2-base-ubuntu22.04 | 87.01 MB | 30.82 MB |
| nvidia/cuda:12.3.2-runtime-ubuntu22.04 | 1.28 GB | 1.23 GB |
| nvidia/cuda:12.3.2-devel-ubuntu22.04 | 3.68 GB | 3.14 GB |
| nvidia/cuda:12.3.2-cudnn9-runtime-ubuntu22.04 | 1.91 GB | 1.86 GB |
| nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04 | 4.31 GB | 3.77 GB |
| nvidia/cuda:12.3.2-base-ubuntu20.04 | 88.25 MB | 32.56 MB |
| nvidia/cuda:12.3.2-runtime-ubuntu20.04 | 1.28 GB | 1.23 GB |
| nvidia/cuda:12.3.2-devel-ubuntu20.04 | 3.66 GB | 3.12 GB |
| nvidia/cuda:12.3.2-cudnn9-runtime-ubuntu20.04 | 1.91 GB | 1.86 GB |
| nvidia/cuda:12.3.2-cudnn9-devel-ubuntu20.04 | 4.29 GB | 3.75 GB |
3. Pytorch 镜像
Pytorch 官方提供的 Pytorch 镜像。
| 镜像 | AMD64 镜像大小 |
|---|
| pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime | 3.44 GB |
| pytorch/pytorch:2.2.2-cuda12.1-cudnn8-devel | 8.47 GB |
4. NVIDIA Pytorch 镜像
Nvidia 打包的 Pytorch 镜像，包含 cuda、cuBlas、cuDNN、OpenMPI、TensorRT 等包。
5. 运行镜像
1
2
3
| docker run --gpus all \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
1
2
3
| docker run --gpus 3 \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
1
2
3
| docker run --gpus '"device=0,1"' \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
除了使用 GPU 的编号，还可以使用 UUID。UUID 可以通过 nvidia-smi -L 查询。
1
2
3
4
| docker run --gpus all \
--shm-size=64g \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
1
2
3
4
| docker run --gpus all \
--ulimit memlock=-1 --ulimit stack=67108864 \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
1
2
3
4
| docker run --gpus all \
--ipc=host \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
1
2
3
4
| docker run --gpus all \
-v "$PWD":/workspace \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
1
2
3
4
| docker run --gpus all \
--cpus=2 --memory=4g \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
1
2
3
4
| docker run --gpus all \
--privileged \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
值得注意的是，Docker 特权模式启动的容器，能使用所有 GPU 资源，导致 --gpus 指定的设备列表无效。
1
2
3
4
5
6
7
| docker run --gpus all \
--ulimit memlock=-1 --ulimit stack=67108864 \
-v "$PWD":/workspace \
--privileged \
--shm-size=64g \
-it --rm \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime bash
|
