使用 NVIDIA Container Toolkit 和 Docker 配置 GPU 深度学习环境

在深度学习开发中，充分利用 GPU 资源是提高模型训练效率的关键。本文将详细介绍如何安装 NVIDIA Container Toolkit，并配置 Docker 容器以利用 GPU 进行深度学习训练。我们将结合 NVIDIA Container Toolkit 的安装和 Docker 容器的配置，确保你拥有一个高效的 GPU 开发环境。

1. 安装 NVIDIA Container Toolkit

NVIDIA Container Toolkit 是 NVIDIA 为 Docker 提供的一个插件，允许容器充分利用 GPU 加速。以下步骤将指导你如何使用 apt 包管理器来安装它。

1.1 配置生产仓库

首先，添加 NVIDIA Container Toolkit 的仓库密钥和源列表：

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

1.2 可选：启用实验性包

如果你希望安装 NVIDIA Container Toolkit 的实验性功能，可以运行以下命令：

sed -i -e '/experimental/ s/^#//g' /etc/apt/sources.list.d/nvidia-container-toolkit.list

1.3 更新包列表并安装 NVIDIA Container Toolkit

更新系统包列表，然后安装 NVIDIA Container Toolkit：

sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit

1.4 重启 Docker 服务

安装完成后，必须重启 Docker 服务以使更改生效：

sudo systemctl restart docker

1.5 验证 NVIDIA Container Toolkit 安装

执行以下命令以验证 NVIDIA Container Toolkit 是否正确安装：

nvidia-container-cli --version

你还可以通过运行以下命令来确认 Docker 是否能够访问 GPU：

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

如果成功显示 GPU 信息，说明 NVIDIA Container Toolkit 已正确安装。

2. 配置 Docker 容器进行 GPU 加速

现在 NVIDIA Container Toolkit 已经安装完毕，我们可以配置 Docker 容器来充分利用 GPU 资源。以下命令将创建一个名为 dev_cuda124 的容器，用于 GPU 深度学习任务：

docker run -it --gpus all \
    -v /mnt:/mnt \
    -v /media:/media \
    --name dev_cuda124 \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    --ipc=host \
    nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04

命令参数解析：

• -it：以交互模式运行容器，方便进行操作。
• --gpus all：允许容器使用主机上的所有 GPU。
• -v /mnt:/mnt 和 -v /media:/media：将主机的 /mnt 和 /media 目录挂载到容器中，以便在容器中访问主机数据。
• --name dev_cuda124：为容器指定名称 dev_cuda124，方便管理和识别。
• --ulimit memlock=-1：取消内存锁定限制，确保程序可以锁定尽可能多的内存，适用于深度学习任务。
• --ulimit stack=67108864：设置线程的最大栈大小为 64MB，避免深度学习任务中出现栈溢出问题。
• --ipc=host：共享主机的 IPC 命名空间，允许容器使用主机的共享内存。

3. 完整配置总结

通过以上步骤，你已经成功配置了 NVIDIA Container Toolkit，并创建了一个可以利用 GPU 进行深度学习任务的 Docker 容器。使用 --ipc=host 确保了容器可以共享主机的共享内存，同时通过 --ulimit 参数确保容器具有足够的栈和内存锁定权限，这将提升深度学习训练的效率。

4. 常见问题与解决方法

• 问题：容器无法识别 GPU

  • 请确保 NVIDIA 驱动和 NVIDIA Container ToolNVIDIA Container Toolkitkit 已正确安装，并通过 nvidia-smi 验证主机是否识别到 GPU。

• 问题：共享内存不足导致数据加载错误

  • 使用 --ipc=host 已经允许容器共享主机的共享内存，但你仍然可以调整主机的 /dev/shm 大小，以满足更大的内存需求。