Docker 与 NVIDIA Container Toolkit
编写时间：2024-04
记录一下 Docker 和 NVIDIA Container Toolkit 各自干什么、怎么配合。以 YOLOv5 训练环境为例。
Docker 和 NVIDIA Container Toolkit 的区别
这俩经常一起出现，但角色完全不同：
| 组件 | 作用 |
|---|---|
| Docker | 容器化运行环境，把项目和依赖打包 |
| NVIDIA Driver | 宿主机上的 GPU 驱动 |
| NVIDIA Container Toolkit | 让容器访问宿主机 GPU |
| CUDA 镜像 | 容器内 CUDA/cuDNN/PyTorch 基础环境 |
Docker 本身不知道 GPU 是什么。NVIDIA Container Toolkit 的角色就是在 docker run 和 GPU 驱动之间架一座桥——它把 GPU 设备文件暴露给容器，让容器里的 CUDA 程序能正常调用驱动。
安装和验证
装完 Docker 后把用户加入 docker 组：
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker验证：
docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi能打印出 nvidia-smi 的输出就说明通了。
YOLOv5 容器示例
搭一个可复用的 YOLOv5 训练镜像，项目结构：
/my-yolov5/
├── Dockerfile
├── docker-compose.yml
├── requirements.txt
├── fonts/
└── yolov5-im/
├── train.py
├── models/
└── data/Dockerfile 基于 PyTorch 官方 CUDA 镜像，装好依赖、复制项目代码即可。这里不展开了。
docker-compose.yml
用 compose 比直接 docker run 更干净——挂载、网络、runtime 都写在一处：
services:
yolov5-im:
build: .
image: yolov5-im:latest
container_name: yolov5_im_train
volumes:
- /path/to/datasets:/workspace/yolov5-im/datasets
working_dir: /workspace/yolov5-im
runtime: nvidia
shm_size: '8gb'
command: sleep infinity
networks:
- yolov5net
networks:
yolov5net:几个值得注意的点：
runtime: nvidia——告诉 compose 用 NVIDIA Container Toolkit，不然容器里看不到 GPUshm_size: '8gb'——PyTorch DataLoader 多 worker 加载时共享内存需求大，默认 64MB 不够command: sleep infinity——让容器保持运行，训练命令手动docker exec进去执行，方便调试
常用操作
docker compose up -d # 启动
docker exec -it yolov5_im_train /bin/bash # 进入容器
docker logs --tail 100 yolov5_im_train # 查看日志
docker cp yolov5_im_train:/workspace/runs ./runs # 复制训练结果
docker compose down # 停止并删除