はじめに
本記事では，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のインストール
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を選択します．
FROM tensorflow/tensorflow:2.7.0-gpu
コンテナ内での権限をローカルと同じにする
参考：dockerでvolumeをマウントしたときのファイルのowner問題
こちらの記事が大変参考になりましたので，詳しい解説は元の記事を参照してください．
これを行わないとDocker内で作成したデータセットや学習済みモデルがroot権限で保存され，
コンテナを抜けた後，ローカルで開けなくなってしまいます．
コンテナ内での権限をローカルと同じにするには，
以下のような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で以下のように記述します．
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でインストールします．
基のイメージのライブラリ一覧を予め確認できればいいのですが，
自分は確認方法がいまいち分からないので，コンテナを起動してから足りない分を追加していってます．
なので一番最初は何もインストールしなくても大丈夫です．
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の全内容
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ファイルを実行するようにしています．
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を動かすために，
環境構築手順と動かし方を説明しました．
自分も動かすまでに様々な記事を参考にさせていただきました．
どれも詳しいかつ分かりやすい記事で大変勉強になりました．
詳しい理解のためにぜひ参考記事をご参照ください．
また，本記事が初投稿となりますので，何かと不備があると思います．
その際はご指摘いただけると幸いです．
何卒よろしくお願いいたします．