GPU マシンを Google Colaboratory のランタイムとして利用する #Docker

せっかく買った GPU マシン¹を有効活用するシリーズ (？)

Google Colab 概要

Jupyter Notebook の Google クローン。

Jupyter Notebook との違いとしては、

Notebook は Google Drive 上に保存される
- 作業内容の共有がしやすい
Google の計算リソースを使える
- GPU, TPU も使える

このあたりが大きい。

Google Colab の良くない点

Google Colab は手っ取り早くデータ分析 / 機械学習ができるのでとても便利なのだけど、素の Jupyter Notebook と比べて微妙な点もいくつかある。

90分 / 12時間ルール

Google Colab でホスト型ランタイムを利用する場合、以下のような制限がある。

何も操作せずに90分経つと自動的に Notebook が終了する
- メモリに載っていた実行結果などが消える
ランタイムのインスタンスが起動してから12時間経つと再起動する
- 追加でインストールしたモジュールなどが消える

これはおそらく「Google のリソースを確保し続けられると困る」という意図の制限なので仕方ないのだけど、「ランチ休憩に行った」とか「すこし別の作業をしていた」とかことあるごとに Notebook のセッションが切れてセルを再実行する羽目になるのでかなり面倒。

ファイルの扱いが面倒

上述の12時間ルールがあるため、Google Colab のホスト型ランタイムではインスタンスのストレージを当てにしてはいけない。

ファイルからデータを読み込んだりデータを永続化したりしたい場合は、Google Drive をランタイムのインスタンスにマウントするなどして対応する必要がある。

補完が遅い

Google のランタイムのインスタンスが物理的に離れているためか、補完が素の Jupyter よりも遅いように感じる。

日本語フォントが入っていない

Matplotlib などで日本語を含むグラフを描画すると豆腐になってしまう。
頑張ればランタイムに日本語フォントをインストールできないことはないが、ランタイムが切り替わるたびにインストール作業をするのが面倒。

ローカルランタイムを利用する

以上の問題は、ローカルランタイムを利用することで解決できる。

90分/12時間ルール → ローカルランタイムではこの制限はない
ファイルの扱いが面倒 → ランタイムのストレージをそのまま使えば問題ない
補完が遅い → 物理的に近くにあるマシンを使うので遅くない
日本語フォントが入っていない → 一度インストールしておけば消えない

ローカルランタイムで GPU マシンを使う

冒頭で言ったとおり「せっかく買った GPU マシン」があるので、Anaconda 3 ベースで且つ TensorFlow GPU が利用できるような Google Colab ランタイムを作ってみることにする。

GPU マシンには NVIDIA ドライバと Docker と NVIDIA Docker 2 をインストールしておく必要があるが、それらのインストール方法は別に書いてあるので割愛。

Ubuntu 18.04 に NVIDIA Docker 環境構築 (2018年秋) - Qiita

Docker イメージを作る

以下のような Dockerfile を書いてイメージをビルドする。

Dockerfile

FROM continuumio/anaconda3

# conda で tensorflow-gpu を入れると CUDA と cuDNN も同時にインストールされる
RUN conda install -y tensorflow-gpu

# NVIDIA container runtime を利用するために必要な設定
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility

# 日本語フォントをインストール
RUN apt update \
  && apt install -y fonts-takao-gothic \
  && apt clean \
  && rm -rf /var/lib/apt/lists/

# matplotlib で日本語フォントを利用するように設定
RUN echo "\nfont.family: TakaoPGothic" >> $(python -c 'import matplotlib as m; print(m.matplotlib_fname())') \
  && rm -f ~/.cache/matplotlib/font*

# Google Colab から接続できるようにするための extension を追加
RUN pip install jupyter_http_over_ws \
  && jupyter serverextension enable --py jupyter_http_over_ws

$ docker build . -t anaconda3-ja:tensorflow-gpu

Jupyter を起動

先程ビルドした Docker イメージを使って Jupyter Notebook を起動する。

--NotebookApp.allow_origin で Google Colab の URL を指定する必要がある。

$ ddocker run -d \
        --runtime=nvidia \
        -p 8888:8888 \
        anaconda3-ja:tensorflow-gpu \
                jupyter notebook \
                        --allow-root \
                        --ip=0.0.0.0 \
                        --NotebookApp.token='' \
                        --NotebookApp.allow_origin='https://colab.research.google.com'

ポートフォワーディングする

手元のノート PC から GPU マシンに対してポートフォワーディングする。

$ ssh ${GPU_HOSTNAME} -N -L 8888:localhost:8888

これで手元のノート PC の 8888 ポートは GPU マシン内の localhost:8888 につながるようになった。

[追記] あまりないユースケースと思うが、SSH できない場合は socat でもできる。

Google Colab からローカルランタイムとして接続する

Google Colab から「ローカルランタイムに接続」を選び、ポート 8888 を指定する。

無事接続できれば完了。

機械学習用の GPU マシンを組んだ - ほとラボ ↩

GPU マシンを Google Colaboratory のランタイムとして利用する