DockerでGPU学習環境構築
背景
ディープラーニングでローカルPCのGPUを使った学習環境を構築した経験のある人は、一度はNVIDIAのドライバやCUDA周りでハマった経験があるのではないでしょうか？そんなバッドノウハウ（怪文章？）をまとめたQiita記事（TensorFlowでGPU学習させるためにCUDA周りではまったときの対処法）に、なんとNVIDIAの中の人(@ksasaki さん）から「Dockerを使えば…人類は幸せになれる（超意訳）」とのコメントをいただきました！
喜び勇んで、NVIDIAのドライバをアップデートしたところ、そこには文鎮と化した起動しないLinuxマシンが…からあげのNVIDIAとの戦いは始まったばかりだ！（戦ってません）
DockerでGPU学習環境構築するメリット
うまく構築できればという前提で、以下のようなメリットがあります。
- 様々なフレームワーク（TensorFlow/PyTorch）、バージョンの使い捨ての学習環境を一瞬で構築できる
- クラウド環境での学習環境構築にもそのまま使える
Docker自体の説明や、メリットに関しては以下記事を参照ください。
DockerでのGPU環境はNVIDIA Container Toolkit（旧 NVIDIA-Docker）を使って実現します。NVIDIA Container Toolkitの概要図は以下となります。ホストPC側は、NVIDIAのGPUドライバさえ最新にしておけば良く、CUDAやTensorFlow/PyTorchのバージョンはDocker側で管理できます。
NVIDIA/nvidia-docker(GitHub)より引用
これでグッと環境構築が楽になる、と言いたいところなのですが、実際やってみるとハマりどころは結構あるもので、2桁回くらいUbuntuを再セットアップすることになりました。その過程で得た、闇の知識をここで惜しげも無く公開したいと思います。
この記事の前提（必要なもの）
- 基礎的なLinuxコマンドの知識
- Ubuntuをインストールするマシン
- USBメモリ
- Ubuntuのバージョンは22.04/20.04/18.04を想定しています
- 黒画やログインループでめげない強い心
使用したハードウェアは、HP ZBOOK、GPUはQuadro P600。デスクトップPC（PC工房のゲーミングPC） + RTX3060等です。
DockerでのGPU学習環境構築の流れ
以下の6つを実施すればOKです。
- Ubuntuインストール
- セキュアブートの無効化
- nouveauの停止
- NVIDIAドライバのインストール
- NVIDIAドライバの確認
- DockerとNVIDIA Container Toolkit（旧 NVIDIA-Docker）のセットアップ
Dockerを使うことで、CUDA、cuDNN、TensorFlowのバージョンの組み合わせに悩まなくて済むのが良いですね。
順に説明していきます。
Ubuntuインストール
Ubuntu 18.04/20.04/22.04を前提としています。UbuntuはUbuntu公式サイトのリリースページから必要なイメージファイルをダウンロードします。
ダウンロードしたイメージファイルをbalenaEtcherというソフトを使用してUSBメモリに書き込みます。balenaEtcherのサイトからダウンロードしてインストールします。
USBメモリをPCに繋いだ状態でbalenaEtherを起動して、以下のようにイメージを選択、USBメモリを選択、Flashをクリックするだけです。
イメージファイルの書き込みには、Raspberry Pi Imagerも使えます。ラズパイになれている人にはこちらもおすすめです。Raspberry Pi Imagerに関しては以下記事参照ください。
USBメモリからのUbuntuのセットアップ方法は省略します（調べたら色々情報出てくると思います）。
使用するUSBメモリ、ドライブに関しては何でもよいですが、一例としては、以下記事参照ください。
セキュアブートの無効化
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
...
Secure bootを無効化できないという闇深い方は、以下記事を参考にすると良いかもしれません（私は未実施です）。
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以降は以下コマンドで自動でセットアップできます。ただ、autoinstallは経験上、黒画になったり、ネットワークが突然死んだりとトラブル多いのでオススメしません。この後紹介する、nvidiaのドライバを手動でインストールする方法がオススメです。
忠告しましたからね！覚悟はできているという人は、以下実行してください。
$ sudo ubuntu-drivers autoinstall
$ sudo shutdown -r now
古いUbuntuの場合は、ubuntu-driversがありませんが、以下コマンドでubuntu-driversをインストールできると思います（未確認）。
$ sudo apt install ubuntu-drivers-common
これでうまくいかない場合（黒画など）は、手動でインストールを試してみましょう。
手動（apt）でインストール
最初にインストールするべきドライバを確認します。以下のコマンドで推奨のドライバがリストアップされます。
$ ubuntu-drivers devices
以下は表示例です。recommendedとあるやつが推奨ドライバのようですが、recommendedを入れても爆死することはよくあります（NVIDIAあるあるです）。
== /sys/devices/pci0000:00/0000:00:01.0/0000:01:00.0 ==
modalias : pci:v000010DEd00002504sv00001462sd0000397Dbc03sc00i00
vendor : NVIDIA Corporation
model : GA106 [GeForce RTX 3060 Lite Hash Rate]
driver : nvidia-driver-525 - distro non-free
driver : nvidia-driver-515-server - distro non-free
driver : nvidia-driver-525-open - distro non-free recommended
driver : nvidia-driver-525-server - distro non-free
driver : nvidia-driver-510 - distro non-free
driver : nvidia-driver-515-open - distro non-free
driver : nvidia-driver-470 - distro non-free
driver : nvidia-driver-470-server - distro non-free
driver : nvidia-driver-515 - distro non-free
driver : xserver-xorg-video-nouveau - distro free builtin
個人的経験からいうと、openとついているやつはrecommendedされても動かないことが多いので、何もついていない無印を選びましょう。
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
PPAパッケージを追加すると、インストールできるドライバが増えているので、再度自分のセットアップしたいバージョンを指定して、aptでインストールします。
$ sudo apt install nvidia-driver-xxx
ドライバのインストールが終わったら再起動します。祈るように起動を待ちましょう。
$ 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のバージョンに落としたり、そのあと455のバージョンで正常に動くようになったりしました。（NVIDIAドライバではよくあることです）。現在（2023年4月）は525バージョンで落ち着いています。
また、ここで表示されているCUDAバージョンは、NVIDIAドライバが対応しているCUDAのバージョンを示しているだけで、実際のCUDAのバージョンとは異なることに注意しましょう。
Sat Apr 1 11:00:09 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.89.02 Driver Version: 525.89.02 CUDA Version: 12.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:01:00.0 On | N/A |
| 0% 39C P5 13W / 170W | 769MiB / 12288MiB | 3% 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 267MiB |
| 0 N/A N/A 1346 G ...ome-remote-desktop-daemon 2MiB |
| 0 N/A N/A 1386 G /usr/bin/gnome-shell 218MiB |
| 0 N/A N/A 2129 G ...428282054872054337,131072 159MiB |
| 0 N/A N/A 2636 G ...RendererForSitePerProcess 118MiB |
+-----------------------------------------------------------------------------+
今回は、Docker使うためCUDAのインストールは実施しませんが、CUDAの正確なバージョンを確認する方法知りたい方は、以下記事を参照ください。
DockerとNVIDIA Container Toolkit(NVIDIA Docker）のインストール
DockerとDockerでGPUを使うために必要なNVIDIA Dockerのインストール方法を記載します。
NVIDIA DockerはNVIDIA Container Toolkitが推奨になったり、またNVIDIA Dockerに戻ったりと混沌としているのですが、2023年4月時点はNVIDIA Dockerが推奨のはずです。NVIDIA Container Toolkitのインストール情報は参考に残しておきます。詳しくはNVIDIA Docker って今どうなってるの？ (20.09 版)を参照ください。
少なくとも2023年7月以降は、また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-Docker
公式サイトの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 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でのGPU動作確認
TensorFlow
Doker上でのGPUの動作確認します。
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 |
+-------------------------------+----------------------+----------------------+
PyTorch
以下のようなDockerファイルを作成します。具体的にはDockerfileというファイル名に以下の内容を書いて保存してください。
FROM ubuntu:20.04
RUN apt-get update && apt-get upgrade -y
RUN apt-get install -y git python3 python3-pip
RUN pip install torch
Dockerビルドします。Dockerfileと同じディレクトリで以下コマンド実行します。
$ docker build -t ubuntu:dockertest .
ビルドしたイメージを走らせます。
$ docker run -it --rm ubuntu:dockertest
コンテナ内で以下実行します。torch.cuda.is_available()がTrueになればOKです。
/# python3
Python 3.8.10 (default, May 26 2023, 14:05:08)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.cuda.is_available())
True
追記：nvidiaのイメージリストは以下にもあります。
A full list of tags that are supported with CUDA, cuDNN, and security updates.
1行目のFROM ubuntu:20.04を以下にしてもPyTorchでCUDAを動かせます。
FROM nvidia/cuda:11.7.1-base-ubuntu22.04
FROM ubuntu:20.04だとCUDAがどこにあるのか分かっていない…
NVIDIAの提供しているNGCイメージ
NVIDAの提供しているNGCのTensorFlow イメージが良いという話を聞いたので試してみました。
tensorflow-release-notesをみて、NVIDIAドライバのバージョンが対応しているイメージを探します（Driver Requirementsのところに書かれています）。私がインストールしている NVIDIAのドライバ440.95.01の場合だと、20.03が対応していたので、以下の通り実行してみます。
$ docker run --gpus all --rm --name ngc -p 8888:8888 -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 Composeのインストール
Docker Composeが必要な場合は、以下コマンドを実行しましょう。
$ sudo apt-get update
$ sudo apt-get install docker-compose-plugin
参考：Install on Linux(Docker Document)
Docker 内 で GPUが使えない (Failed to initialize NVML: Unknown Error) 場合
以下で対応しました。
Dockerで torch.cuda.is_available()がTrueにならない
NVIDIAドライバが古い可能性があるのでアップデートしましょう。以下手順メモです。
まとめ
DockerでディープラーニングのGPU学習環境構築する方法を書きました。色々ハマってしまったので無駄に長くなってしまいました。
Dockerは、数年前から「便利そうだから使ってみよう！」 → 「やっぱり自分には早かった」を繰り返しているような気がしますが、使いこなせれば、CUDAとTensorFlowのバージョン違いに悩まされなくなったり、色々な環境で素早く試したりができたりするので、なんとか使いこなしていきたいなと思います。
関連ページ
参考リンク
- NVIDIA Docker って今どうなってるの？ (20.09 版)
- NVIDIA Docker って今どうなってるの？ (19.11版)
- nvidia-dockerを使用してGPU環境構築
- Ubuntuの18.04 - NVIDIAドライバをインストールする方法
- ubuntu18.04でTensorFlow-gpuを動かす
- Ubuntu Linux 18.04にGPUドライバをインストールする
- UEFI のセキュアブート機にNVIDIAのドライバを入れる話
- 初めてのPCにNVIDIA+Ubuntu+Docker+Tensorflow+ROS2等を入れて、なるべくMacっぽくする
- Ubuntu 20.04 LTS セットアップログ（主に自分用）
変更履歴
- 2023/04/01 Ubuntu22.04を再インストールしたので追記
- 2020/09/09 NVIDIA-Docker2が推奨になったことを追記
- 2020/08/08 コメントのアドバイスを受けて、NVIDIA DockerとNGCイメージに関して文章修正