Docker + PyTorchの環境構築

この記事について

当記事では，Docker上のJupyter NotebookとPyTorchを用いて，Nvidia製のGPUを利用できる環境構築を目指します．

ホストマシンのOSとしてはWindowsを利用します．

しかし，DockerのバックエンドとしてWSL2上に構築したUbuntu環境を用いるので，ホストOSがLinux(特にUbuntu)の場合でも，ほとんど同様の手順で環境構築可能であると考えられます．

環境

• GPU: Nvidia GeForce RTX 3060 (12GB)
• ホストOS: Windows 11 Home
• Docker: 4.31.1
• Dockerのバックエンド: Ubuntu 24.04 (on WSL2)

やること

1. NVIDIA Driverのインストール (Windows)

まず，Windows上にNvidia純正のGPUドライバがインストールされているか確認し，インストールされていない場合は新たにインストールします．

既にインストールされている場合は，たぶんそれを使えます．

詳細はこちらをご確認ください．

2. NVIDIA Container Toolkitのインストール (WSL2)

次に，WSL2上のUbuntuに，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

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

sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

3. CUDAイメージをPull

つぎに，CUDAを環境を含むDockerイメージをPullします．

このとき，利用したいPyTorchのバージョンに対応するCUDAバージョンのイメージをPullするように注意してください．これは，PyTorch公式ドキュメントで確認できます．

今回は，CUDA 12.1をサポートするバージョンをインストールしようと思います．

イメージは，CUDA | NGC CatalogのTagsタブから探せます．

今回は，12.1.1-devel-ubuntu22.04をPullしたいと思います．

docker pull nvcr.io/nvidia/cuda:12.1.1-devel-ubuntu22.04

4. Dockerfileを作ってイメージをビルド

次に，先ほどビルドしたCUDAイメージをもとに，pipコマンドでJupyterとPyTorchをインストールするDockerfileを書きます．

PyTorchをインストールするコマンドは，先述のPyTorch公式ドキュメントを参照してください．

Dockerfile

FROM nvcr.io/nvidia/cuda:12.1.1-base-ubuntu22.04

RUN apt-get update && apt-get install -y python3-pip

RUN python3 -m pip install --upgrade pip \
    && python3 -m pip install --no-cache-dir \
    torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

RUN python3 -m pip install --no-cache-dir jupyter

RUN apt-get autoremove -y && \
    apt-get clean && \
    rm -rf /usr/local/src/*

続いて，このDockerfileをもとに新しいイメージをビルドします．

イメージの名前はtorch-testにしてみました．

docker build -t torch-test .

docker imagesコマンドで，イメージがビルドされたことを確認しておくとよいでしょう．

5. コンテナを起動

docker run --rm -it --gpus all -p 8888:8888 torch-test \
jupyter notebook --allow-root --no-browser --ip 0.0.0.0

--ip 0.0.0.0がないと，notebookが正常に起動しませんでした．(参考: DockerでPytorch & Jupyter Lab環境を簡単構築する)

標準出力に表示されたURLにアクセスすると，いつものJupyter Notebookの画面になります．

最後に，試しに適当なipynbファイルを作って，torchからGPUを利用できるか確認してみましょう．

test.ipynb

import torch
print(torch.cuda.is_available())
print(torch.cuda.get_device_name())

>>> 2.3.1_cu121
    True
    NVIDIA GeForce RTX 3060

正しくGPUが認識されていることがわかります．