NvidiaのDockerイメージでGPU版Tensorflowを動かす(RTX30xxシリーズ対応)

Tensorflowのインストールで困ったこと、ありませんか?

私は、これまでは特にハマることなしにインストールできていたのですが、最近大学のPCに(現時点では)最新のGPUであるRTX3090を導入してから悲劇が始まりました。少なくとも2021年3月時点では、まだTensorflowの公式pipパッケージがRTX30xxシリーズをサポートできていなかったためです。

さて本題ですが、解決方法を端的に言うと公式pipをあきらめてDocker環境(しかもNvidiaの公式イメージ)に移行するというものです。これまでは公式のpipパッケージ(公式と言ってもtf-nightlyというstableではない最新のパッケージ)でだましだましやっていたのですが、Simulation中にPCがフリーズしてしまうことが多々あり、ついにDockerに挑むことにしたのです。

さて、ほんとに本題です。この環境構築のためには3つのソフトをインストールする必要があります。

  1. Docker本体
  2. Nvidia-Docker
  3. Nvidia GPU Cloud(NGC)のTensorflowのDockerイメージ

1.のDockerについては世の中にTutorialが大量にあるので、検索してみてください。簡単に言うと、PCのOS上に別のOSを動かすことができる仕組みで、PC内で別のPCを動かせるみたいなものです。この仕組みを逆手にとって、TensorflowとGPUが動くようなPC環境(これをDockerイメージと呼びます)をあらかじめ用意しておき、ユーザーはそれを持ってきてDocker上で動かせば環境構築の手間いらずでTensorflow+GPUを動かせる、というのが今回やろうとしていることです。

2.のNvidia-DockerというのはDockerでGPU処理を可能にするためのDocker Pluginです。Dockerに加えて必要になります。

また、3.のNvidia GPU Cloudというのは、Cloudという名前が惑わせるのですが、AWSなどのCloudコンピューティングだけを前提にしたものではなく、普通のPCでも使えるものとなっています。この名前を見て、これは違うなと感じてしまう人も多いのではないでしょうか(私も最初はそうでした・・)。ネーミングの大事さを痛感します。

今回試した環境

以下の環境で試していますが、近い環境ならおそらくだいじょうぶだと思います。

環境1

  • OS: Ubuntu 18.04
  • GPU: RTX2070
  • Nvidia driver: 440.33.01

環境2

  • OS: Ubuntu 20.04
  • GPU: RTX3090
  • Nvidia driver: 455.23.05

前準備

最初に、Nvidiaのドライバだけはインストールしておいてください。ドライバのインストール方法は様々ありますが、このサイトでNvidiaの中の人が簡単な方法を解説しています。ただ、NvidiaのGPUの難しいところは、GPUのタイプ毎に使えるドライバ・CUDA・CuDNNのバージョンに縛りがあることです。Nvidiaのサイトに、GPUアーキテクチャ(RTX3090ならAmpere)毎にサポートされるドライバのバージョンが書いてあります。なお、今回は自分でCUDAをインストールする必要はありませんが、このサイトにサポートされているCUDAのバージョンも書いてあります。またNvidiaの別のサイトにCuDNNのバージョンについてもまとめられているので、別途環境を整える必要がある場合は参考になります。

今回必要なのはドライバのみなので、自分のGPUに適合するドライバをインストールしてください。

Docker本体のインストール

基本的にはDocker公式サイトのインストール方法を見ながらやるのが確実だと思いますが、私が行った手順を書いておきます。

前準備として必要なパッケージをインストールします。

> sudo apt-get update
> sudo apt-get install \
    apt-transport-https \
    ca-certificates \
    curl \
    gnupg

DockerのGPG鍵を追加します。

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

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

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

これでDockerのインストールはいったん完了です。以下のコマンドでDockerが正しくインストールされていることを確認してください。

> sudo docker run hello-world

ハロー的な挨拶とDockerの説明などが表示されればOKです。

ただ、ここでひっかかるのがsudo付けないとダメなの?という点です。そこで、ユーザー権限でも動かせるようにします。公式ではまず「sudo groupadd docker」しろと書いてありますが、私の環境ではすでにdockerグループができていました。そこで以下のステップから実行します。

> sudo usermod -aG docker $USER

これでdockerグループに自分が追加されます。以下のコマンドで変化を有効にします。

> newgrp docker 

この状態でもう一度hello-worldを動かします。

> docker run hello-world

これで動いていればDocker環境の構築は完了です。

Nvidia-Dockerのインストール

こちらも基本的には公式のインストールガイドに従うのがよいと思いますが、私の手順ものせておきます。

Stableレポジトリと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-docker2をインストールします。

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

Dockerを再起動します。

> sudo systemctl restart docker

これでインストールは完了です。以下のコマンドを打つと、nvidia-smiのGPU情報が表示されるはずです。

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

なお、ここでエラーが出る場合は、ドライバのバージョンとCUDAのバージョンがあっていない可能性があります。先ほど紹介したNvidiaのサイトで対応するCUDAバージョンを調べて、それを上記の11.0と置き換えて実行してみてください。

Nvidia GPU Cloud(NGC)のTensorflowのDockerイメージのインストール

これも同様に、基本的にはNvidia公式サイトの指示に従ってほしいのですが、これも正しいバージョンを選ばないと動かないのでバージョン選択方法から書きます。バージョン対応表を見ると、Dockerイメージのバージョンごとに対応するドライバのバージョンが書いてあります。例えば、ドライババージョン「455.23.05」に対応するのはDockerイメージの「20.12」であることが分かります。

ということでそのイメージを選択してインストールを行います。公式サイトのTagsというところにイメージの一覧があるので、ここではTensorflowの2系をインストールしたいので「tf2」がついたタグを見つけて、そのタグの「・・・」の部分を押すとインストールのためのpullコマンド(以下)をコピーできます。

> docker pull nvcr.io/nvidia/tensorflow:20.12-tf2-py3

これでインストールは完了です。以下のコマンドの「local_dir」と「container_dir」のところを自分の環境と入れ替えて実行すると、Tensorflowが実行できる環境のshellに入れます。なお、「local_dir」には自分のPC内のディレクトリ(例えば、/home/hoge/work/data)を指定し、「container_dir」には入った環境内でそのディレクトリをmountするディレクトリ(例えば、/data)を指定します。こうすることで、元のPC内のデータにアクセス可能になっています。

docker run --gpus all -it --rm -v local_dir:container_dir nvcr.io/nvidia/tensorflow:xx.xx-tfx-py3

この環境内で、普通にtensorflowのpythonスクリプトを実行すれば動くはずです!

お疲れさまでした!

最後にDocker環境について

Dockerの使い方はクセがあるというか、使いこなすとかなり便利だと思うのですがなかなかハードルが高いです。この記事内では、docker runする時に–rmオプションを付けているので、常に使いすて状態でDockerコンテナを動かす感じで書いていますが、どう使うのがよいのかはDocker初心者の私としてはよくわかっていません。Webを漁れば情報が死ぬほどでてくるので、気になる方は調べてみてください。