使用 NVIDIA Container Toolkit 和 Docker 配置 GPU 深度學習環境

溺死在幸福里的猪發表於2024-10-02
AIDockerGPU深度學習

在深度學習開發中，充分利用 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 大小，以滿足更大的記憶體需求。

相關文章