Ubuntu で nvidia-docker を使って PyTorch の環境を作りたい

:::note warn
`nvidia-docker`は`nvidia-container-toolkit`へ統合され非推奨になりました．新しく記事を書いたのでそちらを参照してください
:::

https://qiita.com/tf63/items/0b137cc60b2559276e61

---

ML の論文の実装を試していると，複数バージョンの CUDA と PyTorch を使い分ける必要があったりします．ローカルで環境構築すると酷い目に合うので，今回は `nvidia-docker` を使ってDockerで環境構築する方法を紹介します

**ソースコード**

https://github.com/tf63/docker-images

**動作環境**

- Ubuntu 22.04
- NVIDIA GeForce RTX 3060 Ti
- Docker Engine ( Docker Desktop は非推奨です )

**事前知識**

https://qiita.com/tf63/items/c21549ba44224722f301

https://qiita.com/tf63/items/684fe4b818ecd715aed9

## 前準備

まず Docker で GPU を認識するのに必要なパッケージをインストールします

https://matsuand.github.io/docs.docker.jp.onthefly/config/containers/resource_constraints/#gpu

```bash
sudo apt install nvidia-container-runtime

# インストール確認
which nvidia-contaier-runtime-hook

sudo apt install nvidia-docker2
```

Docker デーモンを再起動します

```bash
sudo service docker restart
```

## Docker イメージの作成

作成したい環境の `nvidia-docker` イメージを選びます

[Docker hub](https://hub.docker.com/r/nvidia/cuda/tags) から 必要な CUDA のバージョンのイメージを選びましょう

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3291419/4e75f385-f8dc-610b-affb-2de049c83c11.png)

私は CUDA 11.1 or CUDA 11.7 を使うことが多いので，

- `nvidia/cuda:11.1.1-cudnn8-runtime-ubuntu20.04`
- `nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04`

を選んでみます

選んだイメージをベースに `Dockerfile` を作ります

```Dockerfile
FROM nvidia/cuda:11.1.1-cudnn8-runtime-ubuntu20.04

ENV DEBIAN_FRONTEND=noninteractive

WORKDIR /app

RUN apt update && \
  apt install -y \
  wget \
  bzip2 \
  build-essential \
  git \
  git-lfs \
  curl \
  ca-certificates \
  libsndfile1-dev \
  libgl1 \
  python3.8 \
  python3-pip

COPY docker/pytorch/requirements.txt /app

RUN pip3 install --no-cache-dir -U pip && \
  pip3 install --no-cache-dir -r requirements.txt

RUN pip3 install --no-cache-dir torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
```

## Dockerfile で PyTorch をインストールしてみる

https://download.pytorch.org/whl/torch_stable.html で CUDA のバージョンに合った PyTorch を探します

いくつか例を載せておきます

**CUDA11.1**

`cu111/torch-1.9.0%...`，`cu111/torchvision-0.10.0%...`があるので，インストールは

```Dockerfile
  pip3 install --no-cache-dir torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
```

**CUDA11.7**

`cu117/torch-1.13.0%...`，`cu117/torchvision-0.14.0%...`があるので，インストールは

```Dockerfile
  pip3 install torch==1.13.0+cu117 torchvision==0.14.0+cu117 -f https://download.pytorch.org/whl/torch_stable.html
```

**CUDA11.7, PyTorch 2.0**

`cu117/torch-2.0.0%...`，`cu117/torchvision-0.15.0%...`があるので，インストールは

```Dockerfile
  pip3 install torch==2.0.0+cu117 torchvision==0.15.0+cu117 -f https://download.pytorch.org/whl/torch_stable.html
```

## コンテナに GPU を割り当てる

作成したコンテナ上でGPUを利用するには，あらかじめコンテナにGPUを割り当てる必要があります．そのためには GPU 接続したい Docker サービスで deploy を指定します

https://matsuand.github.io/docs.docker.jp.onthefly/compose/gpu-support/

```compose.yaml
services:
    pytorch_111:
        build:
            context: .
            dockerfile: docker/Dockerfile.cu111
        container_name: 'nvidia-docker-111'
        stdin_open: true
        tty: true
        volumes:
            - ./:/app
        deploy:
            resources:
                reservations:
                    devices:
                        - driver: nvidia
                          count: 1
                          capabilities: [gpu]
```

`docker compose up`で起動できるようになっていると思います

```bash
docker compose up -d
```

GPUが利用できるか確認します
```bash
docker compose exec pytorch_111 nvidia-smi
```

## (おまけ) VSCode にコンテナをアタッチしたい

VSCode の拡張機能 `Remote Development` と `Docker` を使うと，コンテナを VSCode にアタッチできます．つまり，ローカル環境で作業するのと同じようにコンテナ上で作業できます

  <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3291419/c8ff1635-5f8e-71e1-5640-0d22f7996e9b.png" width=70%>
  <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3291419/c9290e3c-c72b-e935-234b-053ead212d68.png" width=70%>

`Remote Development` では VSCode のワークディレクトリに `.devcontainer/devcontainer.json` を作成することでコンテナを VSCode にアタッチします

**devcontainer.json の例**

Docker サービス名，`compose.yaml`へのパスなどを指定します

```devcontainer.json
{
    "name": "PyTorch",
    "dockerComposeFile": "../compose.yaml",
    "service": "pytorch_111",
    "workspaceFolder": "/app"
}
```

`ctrl + shift + P`でコマンドパレットを開いて，`Dev Containers: Open Folder in Container...`を実行します

  <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3291419/d334e419-7e1d-a14f-fce1-6a672b00a245.png" width=70%>

ターミナルを開けばコンテナ上のシェルが起動します

  <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3291419/bbc220f8-054e-1d52-4538-dd328aaac2fc.png" width=70%>

PyTorch で GPU を認識できるか確認します

```shell
python
>>> import torch
>>> print(torch.cuda.get_device_name())
NVIDIA GeForce RTX 3060 Ti
```

コンテナから抜ける場合は，左下の`Dev Container: Python`から`Close Remote Connection`を選択してください

  <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3291419/6e80daa1-3e5e-cf31-3496-2605588f47f0.png">

---

以上

## (おまけ)

**docker run で GPU を割り当てたい**

コンテナを起動するときに `--gpus all` を指定すれば ok です

```bash
# Dockerfileをビルドしてイメージを作成
docker build . -f docker/pytorch/Dockerfile.111 -t pytorch_111:1.0 
# イメージからコンテナを起動
docker run -it --rm --gpus all pytorch_111:1.0 nvidia-smi 
```

**Docker Desktop を使う場合**

私がやったときはエラーを吐いてしまったので，Docker Desktop を完全に削除してから Docker Engine を入れ直しました