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

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

---

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

ソースコード

動作環境

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

事前知識

前準備

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

sudo apt install nvidia-container-runtime

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

sudo apt install nvidia-docker2

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

sudo service docker restart

Docker イメージの作成

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

Docker hub から 必要な CUDA のバージョンのイメージを選びましょう

私は 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 を作ります

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%...があるので，インストールは

  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%...があるので，インストールは

  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%...があるので，インストールは

  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 を指定します

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で起動できるようになっていると思います

docker compose up -d

GPUが利用できるか確認します

docker compose exec pytorch_111 nvidia-smi

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

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

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...を実行します

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

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

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

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

---

以上

(おまけ)

docker run で GPU を割り当てたい

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

# 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 を入れ直しました