2021-03-07 / 最終更新日時 : 2024-04-27 naoki009 AI

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

カテゴリー
AI