DockerでGPU+TensorFlow/PyTorchの環境構築した時のメモ

GPU付きサーバのTensorFlow/PyTorchの環境構築して困ったところを中心にメモ。
Dockerは聞いたことあったけど全く使ったことない人です。

2024-11-20追記

最近はVSCode経由でDev Containersを使ってコンテナを管理すると楽。
特にワークスペースごとに異なるコンテナイメージを紐づけたい場合に便利。

2021-12-17追記

Singularityの方が圧倒的に楽だった。Dockerイメージも流用可能。
Dockerだけならrootlessモードを使ってみたい。

なぜGPUサーバの環境構築にDockerを使うのか

風の噂によるとGPUドライバ・ライブラリ周りの互換性が面倒らしい。
Nvidia Container ToolkitでDockerイメージとして入れると幸せになれると聞いた。

OS・ハードウェア

• OS: Ubuntu 20.04.3
• CPU: AMD Ryzen (Zen 2)
• GPU: Nvidia Geforce (Ampere)

Nvidiaドライバのインストール

サーバのベンダーが入れてくれていたので流用。
Driver Version: 470.86

Dockerのインストール

aptでいれた。

$ docker --version
Docker version 20.10.12, build e91ed57

Nvidia Container Toolkitのインストール

公式で特に困らなかった。

$ nvidia-container-cli -V
cli-version: 1.7.0
lib-version: 1.7.0
# 以下略

Dockerの概要・基本操作

Docker操作でsudoが必要になるのを回避

ユーザをDockerのグループに追加すればsudoは不要。

$ sudo usermod -aG docker ${ユーザ名}

docker runのオプション数削減

Dockerfileとdocker-composeを使えば良いらしい。
イメージは固定で引数部分で色々面倒ならdocker-composeだけで済みそう。

docker-composeのインストール

なぜかうまく行かないと思ったら、
aptでインストールしたdocker-composeが古すぎてうまく行かなかった。
見てたマニュアルはv1.27以降のみの記述。

$ sudo apt install docker-compose
$ docker-compose --version
docker-compose version 1.25.0, build unknown
$ sudo apt purge docker-compose

公式情報から最新版をcurlで入れた。

$ docker-compose --version
docker-compose version 1.29.2, build XXX

Dockerにroot以外のユーザで入りたい

docker-compose.xmlを以下のように用意。
イメージはNvidia Container Toolkitのインストール時に試したnvidia/cuda:11.0-baseを利用。

docker-compose.xml

version: '3'
services:
   test:
      image: nvidia/cuda:11.0-base
      volumes:
         - /etc/passwd:/etc/passwd:ro
         - /etc/group:/etc/group:ro
      user: "${UID}:${GID}"

/etc/passwdとgroupはマウントしないとコンテナから読めなくて怒られる。
また、UIDとGIDは環境変数でないとdocker-composeが読めないっぽい。

試した環境だとUIDはシェル変数だけど環境変数ではなく、GIDはそもそも定義されてなかった。
以下のようにするとうまく行った。最初の2行をbashrc等に書いておく。

$ export UID
$ export GID=$(id -g)
$ docker-compose run test /bin/bash
$ docker-compose down

Nvidia公式TensorFlowイメージでメモリ不足の警告

リリースノートを見てドライバのバージョンに合うもの(ここではnvcr.io/nvidia/tensorflow:21.10-tf2-py3)を使う。

$ docker run --gpus all -it --rm nvcr.io/nvidia/tensorflow:21.10-tf2-py3

================
== TensorFlow ==
================
# (中略)

NOTE: The SHMEM allocation limit is set to the default of 64MB.  This may be
   insufficient for TensorFlow.  NVIDIA recommends the use of the following flags:
   docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 ...

# (略)

共有メモリが足りないらしい。
docker-compose.xmlでオプションを指定。

docker-compose.xml

version: '3'
services:
   tensorflow:
      image: nvcr.io/nvidia/tensorflow:21.10-tf2-py3
      volumes:
         - /etc/passwd:/etc/passwd:ro
         - /etc/group:/etc/group:ro
      user: "${UID}:${GID}"
      ipc: host
      ulimits:
         memlock: -1
         stack: -1
      deploy:
         resources:
            reservations:
               devices:
                  - capabilities: [gpu]

Nvidia公式PyTorchイメージ

ほぼtensorflowと同じ。

検証環境まとめ

• OS: Ubuntu 20.04.3
• CPU: AMD Ryzen (Zen 2; x86_64)
• GPU: Nvidia Geforce (Ampere)
• Nvidiaドライバ: 470.86
• Docker: 20.10.12
• Nvidia Container Toolkit: 1.7.0
• docker-compose: 1.29.2

雑感

• 確かに相当スムーズに環境構築出来た。
  • Dockerの学習コストを恐れてたが環境構築だけなら高度な知識は不要そう。
• 個人サーバだからこの設定で行けるが共用サーバだと困りそう。
  • 非管理者がコンテナ操作するならDockerのrootlessモードがよさそう。
  • ライブラリ導入を楽にするのが目的だしSingularityの方が用途に合ってるかも。
  • 複数人が利用するならDockerでJupyterHubサーバを立ててしまうのも検討したい。
• 個人サーバなら素直にライブラリ群入れても良さそうな。
  • Pythonライブラリ間の依存性解消だけを目的とするならconda/pipで仮想環境でも良さそう。
  • Python以外でもGPU周りの環境構築が必要だったので学習コストを払ってもコンテナで運用中。