はじめに

本記事では，DockerでGPUを使ってTensorFlowやPyTorchを動かすために，
環境構築手順と動かし方を説明します．

環境構築

基本的には公式のインストール手順通りで，ところどころ追加で必要なものを入れていきます．
最新バージョンとの違いなどがあるかもしれませんので，参考ページと差異がある場合は参考ページに従ってください．
動作確認はUbuntu18.04で行いました．

NVIDIAドライバの更新

参考：ubuntuにCUDA、nvidiaドライバをインストールするメモ

現在入っているドライバを削除します．

$ sudo apt-get --purge remove nvidia-*
$ sudo apt-get --purge remove cuda-*

NVIDIAドライバのインストールをします．

$ sudo add-apt-repository ppa:graphics-drivers/ppa
$ sudo apt update
$ ubuntu-drivers devices
== /sys/devices/pci0000:00/0000:00:00.0/0000:00:00.0 ==
modalias : pci:0000000000000000000000000000000000000000000000
vendor   : NVIDIA Corporation
model    : XX000 [GeForce XTX 0000 10GB]
driver   : nvidia-driver-390 - distro non-free
driver   : nvidia-driver-470 - distro non-free recommended            ←recommendedされているものを入れます
driver   : xserver-xorg-video-nouveau - distro free builtin
$ sudo apt install nvidia-driver-470                                  ←recommendedされているものを入れます(適宜入れるバージョンを変えてください)
$ sudo reboot

再起動後に以下のコマンドを打って似たような表示が出ればOKです．

$ nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.00    Driver Version: 470.00    CUDA Version: 11.0           |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... Off   | 00000000:00:00.0  On |                  N/A |
| N/A   50C    P0    30W / 250W |   1000MiB /  9313MiB |     15%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1111      G   /usr/lib/xorg/Xorg               1000MiB |
+-----------------------------------------------------------------------------+

Dockerのインストール

参考：Install Docker Engine on Ubuntu

リポジトリを設定するパッケージをインストールします．

$ sudo apt-get update

$ sudo apt-get install \
  apt-transport-https \
  ca-certificates \
  curl \
  gnupg \
  lsb-release

GPG公開鍵を設定します．

$ curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg

リポジトリを設定します．

$ echo \
  "deb [arch=amd64 signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

Dockerをインストールします．

$ sudo apt-get update
$ sudo apt-get install docker-ce docker-ce-cli containerd.io

以下のコマンドを打って似たような表示が出ればOKです．

$ sudo docker run hello-world

Hello from Docker!
This message shows that your installation appears to be working correctly.

~~中略~~

For more examples and ideas, visit:
 https://docs.docker.com/get-started/

毎回sudoを付けなくていいようにします．

$ sudo groupadd docker
$ sudo usermod -aG docker $USER
$ newgrp docker

以下のコマンドを打って似たような表示が出ればOKです．

$ docker run hello-world

Hello from Docker!
This message shows that your installation appears to be working correctly.

~~中略~~

For more examples and ideas, visit:
 https://docs.docker.com/get-started/

NVIDIA Container Toolkitのインストール

参考：Installation Guide

DockerはそのままではGPUが使えません．
よって，Docker上でGPUを使えるようにするために，NVIDIAが提供しているパッケージをインストールします．

GPG公開鍵とリポジトリを設定します．

$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
  && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
  && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

NVIDIA Container Toolkitをインストールします．

$ sudo apt-get update
$ sudo apt-get install -y nvidia-docker2

Dockerを再起動します．

$ sudo systemctl restart docker

以下のコマンドを打って似たような表示が出ればOKです．

$ docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.00    Driver Version: 470.00    CUDA Version: 11.0           |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... Off   | 00000000:00:00.0  On |                  N/A |
| N/A   50C    P0    30W / 250W |   1000MiB /  9313MiB |     15%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1111      G   /usr/lib/xorg/Xorg               1000MiB |
+-----------------------------------------------------------------------------+

動かし方

作成するファイル

動かせるようにするために以下のtrain.py以外の3つのファイルを作ります．
train.pyはご自身の動かしたいプログラムだと思ってください．
今回はすべて同じフォルダに入れておけば大丈夫です．

$ ls
Dockerfile  entrypoint.sh  launch.sh  train.py

Dockerfileを作る

DockerHubで基となるイメージを探す

TensorFlowがイメージを公開しているのでバージョンを選択してコピペします．
https://hub.docker.com/r/tensorflow/tensorflow/

PyTorchはこちらで公開しています．
https://hub.docker.com/r/pytorch/pytorch

また，NVIDIAからもイメージが公開されています．
何が違うかは分かりませんが，こちらの方が処理が早いので，自分は普段こちらを使っています．
TensorFlowもPyTorchもあります．
https://catalog.ngc.nvidia.com/

今回はDockerHubのTensorFlowを使います．
GPUに対応している2.7.0-gpuを選択します．

Dockerfile

FROM tensorflow/tensorflow:2.7.0-gpu

コンテナ内での権限をローカルと同じにする

参考：dockerでvolumeをマウントしたときのファイルのowner問題
こちらの記事が大変参考になりましたので，詳しい解説は元の記事を参照してください．

これを行わないとDocker内で作成したデータセットや学習済みモデルがroot権限で保存され，
コンテナを抜けた後，ローカルで開けなくなってしまいます．

コンテナ内での権限をローカルと同じにするには，
以下のようなentrypoint.shを用意し，

entrypoint.sh

#!/bin/bash

USER_ID=${LOCAL_UID:-9001}
GROUP_ID=${LOCAL_GID:-9001}

echo "Starting with UID : $USER_ID, GID: $GROUP_ID"
useradd -u $USER_ID -o -m user
groupmod -g $GROUP_ID user
export HOME=/home/user

exec /usr/sbin/gosu user "$@"

Dockerfileで以下のように記述します．

Dockerfile

RUN apt-get update && apt-get upgrade -y && apt-get -y install gosu
COPY entrypoint.sh /usr/local/bin/entrypoint.sh
RUN chmod +x /usr/local/bin/entrypoint.sh
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]

Pythonライブラリを追加

自分が使いたいライブラリが基のイメージに無い場合pipでインストールします．
基のイメージのライブラリ一覧を予め確認できればいいのですが，
自分は確認方法がいまいち分からないので，コンテナを起動してから足りない分を追加していってます．
なので一番最初は何もインストールしなくても大丈夫です．

Dockerfile

RUN pip install --upgrade pip
RUN pip install opencv-python
RUN apt-get install -y libgl1-mesa-dev
RUN pip install matplotlib
RUN pip install colorama
RUN pip install pillow
RUN pip install scikit-learn
RUN pip install pandas

コンテナ内にワークスペースを作成する

自分が作業するディレクトリを作成します．
後にここにローカルのディレクトリをマウントします．

Dockerfile

RUN mkdir /mywork

Dockerfileの全内容

Dockerfile

FROM tensorflow/tensorflow:2.7.0-gpu

RUN apt-get update && apt-get upgrade -y && apt-get -y install gosu
COPY entrypoint.sh /usr/local/bin/entrypoint.sh
RUN chmod +x /usr/local/bin/entrypoint.sh
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]

RUN pip install --upgrade pip
RUN pip install opencv-python
RUN apt-get install -y libgl1-mesa-dev
RUN pip install matplotlib
RUN pip install colorama
RUN pip install pillow
RUN pip install scikit-learn
RUN pip install pandas

RUN mkdir /mywork

イメージをビルドする

以下のコマンドでDockerfileからイメージをビルドすることができます．
コンテナ名とタグ名は任意で設定してください．
後でコンテナを起動するときに使用します．
最後の.を忘れないよう注意しましょう．

$ docker image build -t コンテナ名:タグ名 .

ビルドしたイメージを立ち上げる

以下のコマンドでビルドしたイメージを立ち上げることができます．
しかし，これだけだとGPUが使えなかったり，ローカルのファイルを操作できなかったりするので，
ここにオプションをつけていきます．

$ docker container run --rm -it コンテナ名:タグ名 bash

GPUを使えるようにする

以下のオプションを加えてGPUを使えるようにします．
GPU番号はnvidia-smiコマンドでGPUと書いてある場所にある番号です．
--gpus '"device=0,1"'のようにすることで複数選択もできます．（'と"の順番に注意）
--gpus allとすればすべてのGPUを選択できます．

$ docker container run --rm --gpus "device=GPU番号" -it コンテナ名:タグ名 bash

ローカルのディレクトリをコンテナ内にマウントする

以下のオプションを加えてローカルのディレクトリをコンテナ内にマウントすることで，
ローカルのファイルを操作できるようにします．
マウント元は，これから学習・テストなどの作業を行なう場所で，
マウント先は，コンテナ内の先程Dockerfileで作成したディレクトリです．
-v `pwd`:/myworkとすることで現在地をマウントすることができます．（`はバッククオートであることに注意）

$ docker container run --rm --gpus "device=GPU番号" -v ローカルの作業場所の絶対パス:/mywork -it コンテナ名:タグ名 bash

権限を設定する

以下のオプションを加えてコンテナ内の権限をローカルの権限と同じにします．

$ docker container run --rm --gpus "device=GPU番号" -v ローカルの作業場所の絶対パス:/mywork -e LOCAL_UID=$(id -u $USER) -e LOCAL_GID=$(id -g $USER) -it コンテナ名:タグ名 bash

launch.shにする（おまけ）

自分は毎回長いコマンドを打つのが大変なので以下のようなshファイルにしておき，
そのshファイルを実行するようにしています．

launch.sh

docker container run --rm --gpus "device=GPU番号" -v ローカルの作業場所の絶対パス:/mywork -e LOCAL_UID=$(id -u $USER) -e LOCAL_GID=$(id -g $USER) -it コンテナ名:タグ名 bash

$ sh launch.sh

作業開始！

作業ディレクトリの/myworkに移動して作業を開始できます．

$ cd mywork
$ python train.py

おわりに

本記事では，DockerでGPUを使ってTensorFlowやPyTorchを動かすために，
環境構築手順と動かし方を説明しました．

自分も動かすまでに様々な記事を参考にさせていただきました．
どれも詳しいかつ分かりやすい記事で大変勉強になりました．
詳しい理解のためにぜひ参考記事をご参照ください．

また，本記事が初投稿となりますので，何かと不備があると思います．
その際はご指摘いただけると幸いです．
何卒よろしくお願いいたします．