2021-03-07 / 最終更新日時 : 2021-03-12 naoki009 AI

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

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

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

さて本題ですが、解決方法を端的に言うと公式pipをあきらめてDocker環境（しかもNvidiaの公式イメージ）に移行するというものです。その時点で、私は、（名前は知ってはいましたが）Docker？美味しいの？ぐらいの感覚だったので、Dockerでやる道があるというのを薄々知ってはいたものの、ハードル高そうだなと、避けていたというのもありました。また、全ての作業を一気通貫で説明してくれているサイトが見つからなかったというのも、今までやらなかった理由の一つです。これまでは公式のpipパッケージ（公式と言ってもtf-nightlyというstableではない最新のパッケージ）でだましだましやっていたのですが、Simulation中にPCがフリーズしてしまうことが多々あり、ついにDockerに挑むことにしたのです。

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

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

1.のDockerは、まあDockerです。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を漁れば情報が死ぬほどでてくるので、気になる方は調べてみてください。

カテゴリー
AI