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ドライバがインストールされているか確認し，インストールされていない場合は新たにインストールします．

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

詳細は[こちら](https://qiita.com/yu-nakasatomi/items/623dab1949f28db03cde#nvidia%E3%83%89%E3%83%A9%E3%82%A4%E3%83%90%E3%82%92%E3%82%A4%E3%83%B3%E3%82%B9%E3%83%88%E3%83%BC%E3%83%AB%E3%81%99%E3%82%8B)をご確認ください．

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

次に，WSL2上のUbuntuに，NVIDIA Container Toolkitなるものをインストールします．

詳細は[公式インストールガイド](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)を確認してください．

```bash
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
```

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

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

## 3. CUDAイメージをPull

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

このとき，利用したいPyTorchのバージョンに対応するCUDAバージョンのイメージをPullするように注意してください．これは，[PyTorch公式ドキュメント](https://pytorch.org/get-started/locally/)で確認できます．

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

イメージは，[CUDA | NGC Catalog](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/cuda)のTagsタブから探せます．

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

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

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

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

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

```Dockerfile: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`にしてみました．

```bash
docker build -t torch-test .
```

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

## 5. コンテナを起動

```bash
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環境を簡単構築する](https://qiita.com/yu-nakasatomi/items/623dab1949f28db03cde#%E3%82%B3%E3%83%B3%E3%83%86%E3%83%8A%E3%82%92%E8%B5%B7%E5%8B%95%E3%81%99%E3%82%8B))

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

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

```python: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が認識されていることがわかります．