🐳

Devcontainerで構築する深層学習環境

2024/08/07に公開

Docker

VS Code

機械学習

tech

はじめに

VSCodeのDevcontainer機能を使って、深層学習のための開発環境を作ります。今回はライブラリの管理にPoetryを使用します。

なお、以下がインストールされていることを前提にしています。

Visual Studio Code
Docker
最新のNVIDIAドライバ

NVIDIA Container Toolkitをインストール

コンテナ上でGPUを使用するために、NVIDIA Container Toolkitをインストールします。Windowsの場合はWSLにインストールしてください。Ubuntuの場合は以下のコマンドを実行します。

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

sudo apt-get update

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

VSCodeに拡張機能をインストール

Devcontainerを利用するために、以下の拡張機能をインストールします。

設定ファイルを作成する

ここからは開発するディレクトリに移動して行います。
.devcontainer/devcontainer.jsonを作成します。以下に簡単な例を示します。

.devcontainer/devcontainer.json

{
  "name": "python3",
  "image": "nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04",
  "runArgs": ["--gpus", "all", "--shm-size", "16g"],
  "containerEnv": {
    "TZ": "Asia/Tokyo"
  },
  "features": {
    "ghcr.io/devcontainers/features/python:1": {
      "installTools": false,
      "version": "3.10"
    },
    "ghcr.io/devcontainers-contrib/features/poetry:2": {},
    "ghcr.io/devcontainers/features/common-utils:2": {
      "installZsh": false,
      "installOhMyZsh": false,
      "installOhMyZshConfig": false,
      "upgradePackages": false,
      "username": "vscode",
      "userUid": "automatic",
      "userGid": "automatic"
    }
  },
  "remoteUser": "vscode",
  "customizations": {
    "vscode": {
      "extensions": [
        "ms-toolsai.jupyter",
        "charliermarsh.ruff",
        "tamasfe.even-better-toml",
        "ms-python.python"
      ]
    }
  }
}

設定の説明

{
  "image": "nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04"
}

ここで使用するDockerイメージを指定しています。ホストマシンとコンテナのCUDAバージョンを揃える必要はありません。使用するPytorchなどに合わせて選択してください。もちろんimageを指定する代わりにDockerfileやcompose.yamlを使用できます。筆者はすべてのライブラリをPoetryで管理したいので、PytorchなどのイメージではなくCUDAのイメージを使用しています。Pytorchなどのイメージを使用する場合、Poetryは非ルートユーザでライブラリのグローバルインストールを行えないことに注意してください。

{
  "runArgs": ["--gpus", "all", "--shm-size", "16g"]
}

ここでGPUと、shm-sizeを指定しています。shm-sizeは共有メモリ（/dev/shm）で、デフォルトでは64MBです。GPUでモデルをトレーニングすると足りなくなるので増やしています。各自の状況に合わせて値は変更してください。

{
  "features": {
    "ghcr.io/devcontainers/features/python:1": {
      "installTools": false,
      "version": "3.10"
    },
    "ghcr.io/devcontainers-contrib/features/poetry:2": {},
    "ghcr.io/devcontainers/features/common-utils:2": {
      "installZsh": false,
      "installOhMyZsh": false,
      "installOhMyZshConfig": false,
      "upgradePackages": false,
      "username": "vscode",
      "userUid": "automatic",
      "userGid": "automatic"
    }
  }
}

ここで、Pythonのバージョン指定、Poetryのインストール、非ルートユーザの作成などをDockerfileを書かずに行っています。このあたりは好みに合わせて変更してください。

{
  "customizations": {
    "vscode": {
      "extensions": [
        "ms-toolsai.jupyter",
        "charliermarsh.ruff",
        "tamasfe.even-better-toml",
        "ms-python.python"
      ]
    }
  }
}

Devcontainerで使用する拡張機能を指定しています。

Devcontainerの起動

VSCode左下の><をクリック→「コンテナーで再度開く」をクリックします。

Poetryの設定

今回はパッケージマネージャーとしてPoetryを使用するので、コンテナ内で設定します。

# poetryの初期化
poetry init
# ライブラリのインストール例
poetry add torch torchvision
# 開発用ライブラリのインストール例
poetry add ruff ipykernel --group dev

筆者は作成されたpyproject.tomlにpackage-mode = falseを加えています。

.devcontainer/devcontainer.jsonに以下を加えておくと、起動時にライブラリのインストールまで自動で行えます。

{
  "postStartCommand": "poetry install"
}

これで基本的な開発環境の完成です。

おわりに

今回はDevcontainerを使用した深層学習モデルの開発環境の構築方法を紹介しました。誤りを発見したり、「もっといい方法あるよ」という場合はお気軽にコメントいただければ嬉しいです。

GitHubで編集を提案