Ubuntu 20.04のDockerコンテナでTensorFlow GPUサポートを使用

Ubuntu 20.04 に Docker をインストールし、Docker コンテナ上で TensorFlow の GPU サポートを使用します。Ubuntu 20.04 がインストールされている PC に装着しているのは GeForce GTX 10 シリーズのビデオカードです。

TensorFlow のサイトでも、Docker を使用することが TensorFlow の GPU サポートを有効にする最も簡単な方法であると述べられています。

Docker を使用すると、ホストマシンに必要なのは NVIDIA® GPU ドライバだけになるので、Linux で TensorFlow の GPU サポートを有効にする際の最も簡単な方法となります(NVIDIA® CUDA® ツールキットのインストールは不要です)。

https://www.tensorflow.org/install/docker?hl=ja

環境

私の環境は以下のとおりです。

OS

  • Ubuntu 20.04.2 LTS (5.11.0-27-generic)

Docker

  • Docker 20.10.7

NVIDIA ビデオカード

  • MSI GeForce GTX 1050 2GT LP

NVIDIA ドライバ

  • Driver Version: 470.57.02
  • CUDA Version: 11.4

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

Ubuntu 20.04 への NVIDA ドライバのインストール手順は、別のポストに記述しています。

Docker のインストール

Docker の公式ドキュメントを参考に行います。ここではログインしているユーザーを docker グループに入れ、OS 再起動後に docker コマンドを実行できるようにしています。

sudo apt update
sudo apt install -y \
  apt-transport-https \
  ca-certificates \
  curl \
  gnupg \
  lsb-release
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
sudo apt update
sudo apt install -y docker-ce docker-ce-cli containerd.io
sudo usermod -aG docker $USER
sudo reboot

docker run hello-world
Docker Documentation

Jumpstart your client-side server applications with Docker E…

NVIDIA Container Toolkit (nvidia-docker) のインストール

NVIDIA Container Toolkit の GitHub はこちら。

GitHub

Build and run Docker containers leveraging NVIDIA GPUs - Git…

Ubuntu へのインストールガイドはこちら。

インストールガイドの通り、nvidia-docker2 パッケージをインストールします。

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
sudo apt update
sudo apt install -y nvidia-docker2
sudo systemctl restart docker

NVIDIA CUDA イメージでの動作確認

NVIDIA が提供している CUDA イメージのコンテナで GPU が認識されるか動作を確認します。

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

初回はイメージのダウンロードが発生します。結果として nvidia-smi コマンドの実行結果が表示されることを確認します。

Unable to find image 'nvidia/cuda:11.0-base' locally
11.0-base: Pulling from nvidia/cuda
...
Status: Downloaded newer image for nvidia/cuda:11.0-base
Tue Aug 19 13:49:39 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.57.02    Driver Version: 470.57.02    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| 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 |
| 46%   46C    P0    N/A /  75W |   1125MiB /  1966MiB |     15%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

TensorFlow Docker のイメージでの動作確認

GPU サポートと Jupyter が含まれた最新リリースのイメージをダウンロードし、TensorFlow を実行してみます。

TensorFlow のドキュメントはこちら。

TensorFlow

TensorFlow の DockerHub はこちら。

docker pull tensorflow/tensorflow:latest-gpu-jupyter
docker run --gpus all -it --rm tensorflow/tensorflow:latest-gpu-jupyter \
   python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

実行結果です。

2021-08-19 14:48:40.123068: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.127448: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.127641: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.128040: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-08-19 14:48:40.128516: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.128679: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.128822: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.462004: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.462245: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.462428: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-19 14:48:40.462593: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 516 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1050, pci bus id: 0000:01:00.0, compute capability: 6.1
tf.Tensor(-1737.8235, shape=(), dtype=float32)

コンテナ内でシェルを実行したい場合、

docker run --gpus all -p 8888:8888 -it tensorflow/tensorflow:latest-gpu-jupyter bash

コンテナの起動と同時に Jupyter Notebook を起動したい場合、

docker run --gpus all -p 8888:8888 -it tensorflow/tensorflow:latest-gpu-jupyter