Ubuntu 18.04 に TensorFlow Docker を導入する - 子育てしながらエンジニアしたい

前回の記事で、Ubuntu 18.04 に GPU Docker 環境を導入しました。

edosha.hatenablog.jp

今回はこの Docker 上で TensorFlow Docker を動かしてみたいと思います。

TensorFlow Docker とは

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

TensorFlow Docker の公式サイトには上記のことが書かれています。
これはかなり大きいメリットだと思います。

TensorFlow のバージョンによって CUDA とか CuDNN の必要なバージョンも変わります。
一つの環境でやる分には良いですが、お客さんごとに使っている環境が違う場合、これが致命的になります。
TensorFlow Docker があれば、一つのマシンで複数の CUDA / CuDNN のバージョンを切り替えられることになるので、とってもありがたいのです。

所望のバージョンの TensorFlow Docker を pull する

以下のサイトから所望のバージョンを探します。

hub.docker.com

バージョンや中に入っているパッケージの種類ごとにタグが切られています。
たとえばバージョン 2.0.0 で GPU サポートあり、Python3、jupyter 環境込であれば、

2.0.0-gpu-py3-jupyter

となります。

見つかったら Ubuntu のコンソールで以下のようにコマンドを叩きます。

$ docker pull tensorflow/tensorflow:2.0.0-gpu-py3-jupyter

TensorFlow Docker を起動する

先程の公式サイトでは、GPU を使う方法として、nvidia-container-toolkit の昔のやり方が書かれていました。
でも前回の記事のとおり GPU サポートのコマンドが変わりましたので、こちらでは新しい方法で起動してみます。

上で pull した TensorFlow Docker は、GPU サポートあり、jupyter パッケージ込みのものです。
これをDocker で起動するには以下のようにします。

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

- gpus all のところが、以前と変わった部分です。
p 8888:8888 は、Docker 内の Jupyter サーバーに、Docker 外からアクセスできるようにするためのポート指定です。

このコマンドを叩くと以下のような出力がでました。

________                               _______________                
___  __/__________________________________  ____/__  /________      __
__  /  _  _ \_  __ \_  ___/  __ \_  ___/_  /_   __  /_  __ \_ | /| / /
_  /   /  __/  / / /(__  )/ /_/ /  /   _  __/   _  / / /_/ /_ |/ |/ / 
/_/    \___//_/ /_//____/ \____//_/    /_/      /_/  \____/____/|__/


WARNING: You are running this container as root, which can cause new files in
mounted volumes to be created as the root user on your host machine.

To avoid this, run the container by specifying your user's userid:

$ docker run -u $(id -u):$(id -g) args...

[I 02:37:54.373 NotebookApp] Writing notebook server cookie secret to /root/.local/share/jupyter/runtime/notebook_cookie_secret
jupyter_http_over_ws extension initialized. Listening on /http_over_websocket
[I 02:37:54.485 NotebookApp] Serving notebooks from local directory: /tf
[I 02:37:54.485 NotebookApp] The Jupyter Notebook is running at:
[I 02:37:54.485 NotebookApp] http://6aacb5541856:8888/?token=16cd1462e9d40ba3b91a837c2240f097ab75ed173a3bec14
[I 02:37:54.486 NotebookApp]  or http://127.0.0.1:8888/?token=16cd1462e9d40ba3b91a837c2240f097ab75ed173a3bec14
[I 02:37:54.486 NotebookApp] Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).
[C 02:37:54.488 NotebookApp] 
    
    To access the notebook, open this file in a browser:
        file:///root/.local/share/jupyter/runtime/nbserver-1-open.html
    Or copy and paste one of these URLs:
        http://6aacb5541856:8888/?token=16cd1462e9d40ba3b91a837c2240f097ab75ed173a3bec14
     or http://127.0.0.1:8888/?token=16cd1462e9d40ba3b91a837c2240f097ab75ed173a3bec14

出力された URL に Ubuntu のブラウザーでアクセスすると、おなじみの Jupyter 画面が出力されます。

f:id:edosha:20200218114341p:plain — Jupyter

これだけで、GPU サポート込みの環境が作れるとは、ちょっと驚きの快適さでした...