DockerでGPU学習環境構築

背景

　ディープラーニングでローカルPCのGPUを使った学習環境を構築した経験のある人は、一度はNVIDIAのドライバやCUDA周りでハマった経験があるのではないでしょうか？そんなバッドノウハウ（怪文章？）をまとめたQiita記事（TensorFlowでGPU学習させるためにCUDA周りではまったときの対処法）に、なんとNVIDIAの中の人(@ksasaki さん）から「Dockerを使えば…人類は幸せになれる（超意訳）」とのコメントをいただきました！

　喜び勇んで、NVIDIAのドライバをアップデートしたところ、そこには文鎮と化した起動しないLinuxマシンが…からあげのNVIDIAとの戦いは始まったばかりだ！（戦ってません）

DockerでGPU学習環境構築するメリット

　うまく構築できればという前提で、以下のようなメリットがあります。

様々なフレームワーク（TensorFlow/PyTorch）、バージョンの使い捨ての学習環境を一瞬で構築できる
クラウド環境での学習環境構築にもそのまま使える

　Docker自体の説明や、メリットに関しては以下記事を参照ください。
Docker入門して機械学習環境構築

　ただ、やってみるとハマりどころは結構あるもので、2桁回くらいUbuntuを再セットアップすることになりました。その過程で得た、闇の知識をここで惜しげも無く公開したいと思います。

この記事の前提（必要なもの）

基礎的なLinuxコマンドの知識
Ubuntuをインストールするマシン
USBメモリ
Ubuntuのバージョンは18.04を想定(16.04でも多分OK）
黒画やログインループでめげない強い心

　使用したハードウェアは、HP ZBOOK、GPUはQuadro P600です。最新のいけてるGPUなら、ハマる確率は小さいのではないかと思います。

DockerでのGPU学習環境構築の流れ

　以下の6つを実施すればOKです。

Ubuntuインストール
セキュアブートの無効化
nouveauの停止
NVIDIAドライバのインストール
NVIDIAドライバの確認
DockerとNVIDIA Container Toolkit（旧 NVIDIA-Docker）のセットアップ

　Dockerを使うことで、CUDA、cuDNN、TensorFlowのバージョンの組み合わせに悩まなくて済むのが良いですね。

　順に説明していきます。

Ubuntuインストール

　Ubuntu 18.04を前提としています。UbuntuはUbuntu公式サイトの18.04イメージからイメージファイルをダウンロードします。

　ダウンロードしたイメージファイルをbalenaEtcherというソフトを使用してUSBメモリに書き込みます。balenaEtcherのサイトからダウンロードしてインストールします。

　USBメモリをPCに繋いだ状態でbalenaEtherを起動して、以下のようにイメージを選択、USBメモリを選択、Flashをクリックするだけです。

　USBメモリからのUbuntuのセットアップ方法は省略します（調べたら色々情報出てくると思います）。

セキュアブートの無効化

　BIOSのセキュアブートは無効化しないと、後々NVIDIAのドライバを確認するnvidia-smiコマンドで以下のようなエラーメッセージが出てしまいます。必ず実施しましょう（サボってハマりました）。

NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver. Make sure that the latest NVIDIA driver is installed and running.

　PCによって方法が異なりますので手順は割愛します。BIOSを起動して、Secure BootをDisableにするのが基本です。

　Secureブートが無効化されたかは、Ubuntuのターミナルで以下のコマンドを実行して確認できます。

$ dmesg | grep Secure

　以下のようにSecure boot disabledと表示されたらOKです。

[    0.000000] secureboot: Secure boot disabled
...

nouveauの停止

　NVIDIAのドライバにとって邪魔なグラフィックドライバであるnoveauを停止します。以下でnouveauが動いているか確認します。

$ lsmod | grep -i nouveau

　以下のように表示されたら、nouveauが動いています。

nouveau              1896448  1
mxm_wmi                16384  1 nouveau
ttm                   102400  1 nouveau
drm_kms_helper        184320  2 i915,nouveau
drm                   491520  8 drm_kms_helper,i915,ttm,nouveau
i2c_algo_bit           16384  2 i915,nouveau
wmi                    32768  5 hp_wmi,intel_wmi_thunderbolt,wmi_bmof,mxm_wmi,nouveau
video                  49152  2 i915,nouveau

　止めるために、以下の設定ファイルを作成します。

$ sudo vi /etc/modprobe.d/blacklist-nvidia-nouveau.conf

　blacklist-nvidia-nouveau.confの中身は以下にしてください。

blacklist nouveau
options nouveau modset=0

　以下コマンドで設定を有効します。

$ sudo update-initramfs -u

再起動しないまま次に進みます（NVIDIAドライバをインストールしないまま再起動すると画面が崩れる可能性があります）。

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

　自動でインストールする方法と手動（apt）でインストールする方法の2つがあります。自動が楽で良いですが、ダメだった場合のために、手動の方法も説明します。

　自分の環境では、Ubuntuセットアップ直後で自動でインストールする方法が一番良かったです（他は全部ダメで、Ubuntu再インストールし続けました）。

自動でインストール

　最初に、なるべく新しいドライバをインストールしたい場合は、以下コマンドであらかじめPPAパッケージのリストを追加しましょう。ただしPPAなので自己責任で。

$ sudo add-apt-repository ppa:graphics-drivers/ppa
$ sudo apt update

　Ubutu 18.04以降は以下コマンドで自動でセットアップできます。

$ sudo ubuntu-drivers autoinstall
$ sudo shutdown -r now

　古いUbuntuの場合は、ubuntu-driversがありませんが、以下コマンドでubuntu-driversをインストールできると思います（未確認）。

$ sudo apt install ubuntu-drivers-common

　これでうまくいかない場合（黒画など）は、手動でインストールを試してみましょう。

手動（apt）でインストール

　最初にインストールするべきドライバを確認します。以下のコマンドで推奨のドライバがリストアップされます。

$ ubuntu-drivers devices

以下は表示例です。recommendedとあるやつが推奨ドライバです（多分）。

ubuntu-drivers devices
== /sys/devices/pci0000:00/0000:00:01.0/0000:01:00.0 ==
modalias : pci:v000010DEd00001CBCsv0000103Csd0000847Bbc03sc00i00
vendor   : NVIDIA Corporation
driver   : nvidia-driver-440-server - distro non-free recommended
driver   : nvidia-driver-390 - distro non-free
driver   : nvidia-driver-435 - distro non-free
driver   : nvidia-driver-410 - third-party free
driver   : nvidia-driver-440 - distro non-free
driver   : nvidia-driver-418-server - distro non-free
driver   : xserver-xorg-video-nouveau - distro free builtin

　ubuntu-driversを使えない場合は、自分のGPUの型番を以下コマンドで調べてから、NVIDIAのサイトで対象のドライバを検索しましょう。

$ sudo lshw -C display

　あとはaptでお目当てのドライバをインストールしましょう。xxxには自分のセットアップしたいバージョンを指定しましょう（TABで候補が出てきます）。

$ sudo apt install nvidia-driver-xxx

　より新しいドライバをインストールしたい場合は、自動でインストールする場合と同様、以下コマンドでPPAパッケージのリストを追加しましょう。

$ sudo add-apt-repository ppa:graphics-drivers/ppa
$ sudo apt update

　あとは、再起動します。祈るように起動を待ちましょう。

$ sudo shutdown -r now

　再起動して無事起動画面が表示されたらOKです。おめでとうございます。

　不幸にも黒画になったり、起動しなかった場合は残念でした。頑張って修復するか、もういちどUbuntuインストールから始めましょう。NVIDIAのドライバを変えると起動するようになるかもしれません。

　この後、Dockerを動かすとき、最新のイメージほどホスト側のPCの最新のNVIDIAドライバが必要とされます。NVIDIAのドライバが対応していないと、Dockerを動かしたとき、以下のように「NVIDIAのドライバのバージョンが古いよ！」と怒られてしまいます

ERROR: This container was built for NVIDIA Driver Release 450.51 or later, but
       version 440.95.01 was detected and compatibility mode is UNAVAILABLE.

NVIDIA ドライバの確認

　再起動した後、正しくドライバがインストールされているかを確認します。その前に nouveau が停止しているかを確認しましょう。以下を実行して何も表示されなければOKです。

$ lsmod | grep -i nouveau

続いて以下のコマンドでNVIDIAのドライバ情報を確認します。

$ nvidia-smi

　以下のようにドライバのバージョンが表示されたらOKです。450のバージョンだとXがどうしても表示できないため、しかたなく440のバージョンに落としています。CUDA 10.2が表示されていますが、何かの間違いでしょう（インストールしていないので）。

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.95.01    Driver Version: 440.95.01    CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Quadro P600         Off  | 00000000:01:00.0 Off |                  N/A |
| N/A   57C    P0    N/A /  N/A |    476MiB /  4040MiB |      2%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      1143      G   /usr/lib/xorg/Xorg                           335MiB |
|    0      1427      G   /usr/bin/gnome-shell                         137MiB |
+-----------------------------------------------------------------------------+

CUDAの正確な情報を知りたい場合はnvccを使いましょう。以下はnvccで確認する方法です（nvccでの確認に関しては、飛ばしてもOKです）。

　aptでnvidia-cuda-toolkitをインストールします。

$ sudo apt install nvidia-cuda-toolkit

　あとは、以下コマンドでCUDAの情報を取得できます。

$ nvcc -V

　何故かこちらではCUDA9.1がインストールされていることになっています。謎ですね…

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2017 NVIDIA Corporation
Built on Fri_Nov__3_21:07:56_CDT_2017
Cuda compilation tools, release 9.1, V9.1.85

　気にせず（オイ）次にいきましょう。誰か詳しい人は教えてください…

DockerとNVIDIA Container Toolkit(NVIDIA Docker）のインストール

　DockerとDockerでGPUを使うために必要なNVIDIA Container Toolkitをインストールします。参考までに、少し古いNVIDIA Docker（公式ではNVIDIA Container Toolkitが推奨されています）のセットアップ方法も記載します。

Dockerのインストール

　UbuntuにDockerをインストールします。オフィシャルの手順がやや煩雑なので手軽にセットアップできるスクリプト用意しました。

　最初にcurlをインストールします

$ sudo apt install -y curl

　その後、以下のコマンドを実行すればoKです。

$ curl -s https://raw.githubusercontent.com/karaage0703/ubuntu-setup/master/install-docker.sh | /bin/bash

　ここで、権限を有効にするために、一旦再起動しましょう（再ログインで良いはずなのですが、私の環境では駄目だったため再起動しました）。

$ sudo shutdown -r now

NVIDIA Container Toolkitのインストール

NVIDIAのnvidia-container-toolkitパッケージをインストールします。Quickstartの手順をスクリプト化したので、以下の通り実行すればOKです。

$ curl -s https://raw.githubusercontent.com/karaage0703/ubuntu-setup/master/install-nvidia-container-toolkit.sh | /bin/bash

　続いてDockerを起動するため、以下コマンド実行すると…

$ docker run --gpus all --rm nvidia/cuda nvidia-smi

docker: Error response from daemon: OCI runtime create failed: container_linux.go:349: starting container process caused "process_linux.go:449: container init caused \"process_linux.go:432: running prestart hook 0 caused \\\"error running hook: exit status 1, stdout: , stderr: nvidia-container-cli: requirement error: unsatisfied condition: cuda>=11.0, please update your driver to a newer version, or use an earlier cuda container\\\\n\\\"\"": unknown.

　cudaがどうこう言われてダメでした…NVIDIAのドライバが古いのでしょうか…

　とりあえず、TensorFlow公式のDockerイメージで試してみます。以下コマンドでDockerを動かします。

$ docker run --gpus all -it --rm --name tensorflow-gpu -p 8888:8888 tensorflow/tensorflow:latest-gpu-py3-jupyter

　実行すると勝手にイメージのダウンロードが始まり、ダウンロードが完了するとコンテナが起動します。起動すると画面に token=xxxxxxxxxxxxxx という形でtokenが表示されます。

　続いて、ブラウザで以下のアドレスにアクセスしましょう。tokenの後は、先ほど表示されたtokenを入力します。

http://127.0.0.1:8888/?token=xxxxxxxxxxxxxx

　Jupyter Notebookにアクセスできました。

　以下コマンド実行して、上図のようにGPUが表示されたらOKです。

from tensorflow.python.client import device_lib
device_lib.list_local_devices()

　何か動かしてみたい人は、tensorflow-tutorialというディレクトリの中に、色々なサンプルが入っているのでMNISTあたりを動かしてみると良いと思います。nvidia-smiを起動すれば、メモリが消費されてちゃんとGPUが動いていることが確認できます。

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.95.01    Driver Version: 440.95.01    CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Quadro P600         Off  | 00000000:01:00.0 Off |                  N/A |
| N/A   55C    P0    N/A /  N/A |   3874MiB /  4040MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

NVIDIA-Docker

　参考までにNVIDIA-Dockerのインストール方法を紹介します。繰り返しになりますが、NVIDIA-Dockerは古いバージョンとなり、公式はNVIDIA Container Toolkitを推奨していますので、特別な理由がなければNVIDIA Container Toolkitを利用ください。

　こちらも公式サイトのREADMEをもとに、セットアップスクリプト作成しました。以下コマンドでインストールできます。

$ curl -s https://raw.githubusercontent.com/karaage0703/ubuntu-setup/master/install-nvidia-docker.sh | /bin/bash

　以下のnvidia-dockerコマンドで、同様に動かせます。

$ nvidia-docker run -it --rm --name tensorflow-gpu -p 8888:8888 tensorflow/tensorflow:latest-gpu-py3-jupyter

NVIDIAの提供しているNGCイメージ

　NVIDAの提供しているNGCのTensorFlow イメージが良いという話を聞いたので試してみました。

　tensorflow-release-notesをみて、NVIDIAドライバのバージョンが対応しているイメージを探します（Driver Requirementsのところに書かれています）。私がインストールしている NVIDIAのドライバ440.95.01の場合だと、20.03が対応していたので、以下の通り実行してみます。

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

　一瞬起動画面は出るのですが、Dockerにログインできず…他のバージョンもいくつか試してみましたがダメでした・・・残念ながら諦めました。原因はわかりませんでした。NVIDIAのドライバの関係でしょうか…？

　コメントのアドバイス通り、オプションを修正することで無事Dockerにログインできました！NGCイメージは、TensorFlowはもちろん、TensorRTやDALI、豊富なサンプル、Jupyter Labなどが含まれていて便利そうです。

黒画になってしまった場合

　とりあえずリカバリーモードでログインするのが良いです。起動時に「ESC」連打で入れます。

　「ESC」でブート選択画面になるPCの場合は、Continue boot（ブートを続ける）選択をした直後に「ESC」を1回押すと入れたりします。

　rootでCUI画面に入ったら、以下でドライバを削除すると黒画から脱出できたりします。

$ sudo apt --purge autoremove nvidia*

　どうしてもダメなら、素直に再インストールすれば良いと思います。

まとめ

　DockerでディープラーニングのGPU学習環境構築する方法を書きました。色々ハマってしまったので無駄に長くなってしまいました。

　Dockerは、数年前から「便利そうだから使ってみよう！」 → 「やっぱり自分には早かった」を繰り返しているような気がしますが、使いこなせれば、CUDAとTensorFlowのバージョン違いに悩まされなくなったり、色々な環境で素早く試したりができたりするので、なんとか使いこなしていきたいなと思います。

参考リンク

変更履歴

2020/08/08 コメントのアドバイスを受けて、NVIDIA DockerとNGCイメージに関して文章修正

DockerでのディープラーニングGPU学習環境構築方法