Neo Sun


What's past is prologue

在 Linux 下使用 Docker 配置 NVIDIA 4090 深度学习环境

  • ~4.81K words

完整教程：在 Linux 下使用 Docker 配置 NVIDIA 4090 深度学习环境

第一部分：宿主机环境准备

1. 安装 Linux 系统

  • 推荐发行版：Ubuntu 22.04 LTS（对 NVIDIA 驱动支持最佳）
  • 安装方式：
    • 物理机安装：制作启动盘，彻底替代现有系统。
    • 虚拟机安装（仅测试用途）：需启用 PCIe 直通（对 GPU 性能有损耗）。
  • 系统更新：
    1
    sudo apt update && sudo apt upgrade -y

2. 安装 NVIDIA 驱动

  • 禁用 Nouveau 驱动：
    1
    2
    3
    sudo bash -c "echo 'blacklist nouveau' >> /etc/modprobe.d/blacklist.conf"
    sudo update-initramfs -u
    sudo reboot
  • 安装驱动：
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    # 添加官方驱动源
    sudo add-apt-repository ppa:graphics-drivers/ppa
    sudo apt update

    # 查找适用于 4090 的最新驱动
    ubuntu-drivers devices

    # 安装推荐驱动（例如 535 版本）
    sudo apt install nvidia-driver-535
    sudo reboot
  • 验证安装：
    1
    nvidia-smi  # 应显示 GPU 信息和 CUDA 版本

第二部分：配置 Docker 与 GPU 支持

1. 安装 Docker

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 卸载旧版本
sudo apt remove docker docker-engine docker.io containerd runc

# 设置仓库
sudo apt install -y ca-certificates curl
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
sudo chmod a+r /etc/apt/keyrings/docker.gpg

# 添加源
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
  sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

# 安装 Docker
sudo apt update
sudo apt install -y docker-ce docker-ce-cli containerd.io

# 验证安装
sudo docker run hello-world

2. 配置 NVIDIA Container Toolkit

1
2
3
4
5
6
7
8
9
10
11
12
# 添加源
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-container-toolkit
sudo systemctl restart docker

# 验证 GPU 支持
docker run --gpus all --rm nvidia/cuda:12.2.0-base-ubuntu22.04 nvidia-smi

第三部分：Docker 核心操作指南

1. 镜像管理

命令 说明
docker images 查看本地镜像
docker pull nvidia/cuda:12.2.0-devel-ubuntu22.04 拉取镜像
docker rmi <image_id> 删除镜像

2. 容器生命周期

命令 说明
docker run -it --gpus all ubuntu bash 启动新容器
docker start <container_id> 启动已停止的容器
docker stop <container_id> 停止运行中的容器
docker rm <container_id> 删除容器

3. 数据持久化

1
2
3
4
5
6
# 挂载主机目录到容器
docker run -v /host/path:/container/path ...

# 创建 Docker 卷
docker volume create mydata
docker run -v mydata:/container/path ...

4. 网络配置

1
2
3
4
5
# 映射端口
docker run -p 8080:80 ...

# 查看网络
docker network ls

第四部分：构建深度学习容器

1. 创建 Dockerfile

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# 使用 PyTorch 官方镜像（CUDA 12.1）
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime

# 安装系统依赖
RUN apt update && apt install -y \
    build-essential \
    cmake \
    git \
    wget \
    && rm -rf /var/lib/apt/lists/*

# 安装 Python 工具链
RUN pip install --no-cache-dir \
    setuptools \
    wheel \
    ninja \
    scikit-build

# 设置工作目录
WORKDIR /app
COPY . .

# 编译自定义 C 扩展
RUN cd c_extensions && \
    pip install .

# 设置启动命令
CMD ["python", "train.py"]

2. 构建镜像

1
docker build -t my-dl-project .

3. 运行容器

1
2
3
4
docker run -it --rm --gpus all \
  -v $(pwd)/data:/app/data \
  -v $(pwd)/results:/app/results \
  my-dl-project

第五部分：多 CUDA 版本管理

1. 项目 A：CUDA 12.2 + PyTorch 2.3

1
2
3
# Dockerfile.cuda12
FROM nvidia/cuda:12.2.0-devel-ubuntu22.04
RUN pip install torch==2.3.0 torchvision==0.18.0

2. 项目 B：CUDA 11.8 + TensorFlow 2.13

1
2
3
# Dockerfile.cuda11
FROM nvidia/cuda:11.8.0-devel-ubuntu20.04
RUN pip install tensorflow==2.13.0

3. 启动不同版本容器

1
2
3
4
5
6
7
# 构建镜像
docker build -t project-a -f Dockerfile.cuda12 .
docker build -t project-b -f Dockerfile.cuda11 .

# 运行容器
docker run -it --gpus all project-a
docker run -it --gpus all project-b

第六部分：高级技巧

1. 多阶段构建优化

1
2
3
4
5
6
7
8
9
10
# 构建阶段
FROM nvidia/cuda:12.2.0-devel-ubuntu22.04 as builder
RUN apt install -y cmake
COPY . .
RUN make

# 运行阶段
FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04
COPY --from=builder /app/bin /app
CMD ["/app/main"]

2. 使用 Docker Compose

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# docker-compose.yml
version: '3.8'
services:
  trainer:
    image: my-dl-project
    runtime: nvidia
    volumes:
      - ./data:/app/data
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

第七部分：验证与调试

1. 检查 GPU 是否可用

1
2
3
import torch
print(torch.cuda.is_available())          # 应输出 True
print(torch.cuda.get_device_name(0))      # 应显示 "NVIDIA GeForce RTX 4090"

2. 查看 CUDA 版本

1
2
3
# 容器内执行
nvcc --version
cat /usr/local/cuda/version.json

3. 常见问题排查

  • CUDA 不可用：
    • 检查 --gpus all 参数
    • 确认 nvidia-container-toolkit 已安装
  • 权限问题：
    1
    2
    sudo usermod -aG docker $USER   # 将用户加入 docker 组
    newgrp docker                   # 刷新组权限

总结

通过本教程，你可以：

  1. 在 Linux 系统中为 NVIDIA 4090 配置深度学习环境
  2. 使用 Docker 实现完全隔离的 Python/CUDA 环境
  3. 自由切换不同 CUDA 版本（11.x/12.x）
  4. 编译 C/C++ 扩展并打包为 wheel
  5. 通过数据卷持久化实验结果
This work is licensed under Attribution-ShareAlike 4.0 International (CC BY-SA 4.0)