機械学習を使った研究をする環境をDocker + JupyterLab + TensorBoard + cookiecutterで作ってる話
概要
- PyTorchでモデルを書いてGPUぶん回す日々を過ごしているが，複数のタスクを並行でやったりするので環境の切り分けやコード・データの管理が大変だったのをなんとかしようとした
- Dockerの導入とTensorBoardで学習の経過を可視化できるようにしたり，色んな工夫をやってたら大分マシになった
モチベーション
研究でプログラムを書いてると，実行環境やコード・データの管理にかかる手間って意外とでかい． 効率化のためにいろいろ工夫してきたけど，やってることが増えすぎて全部頭に入れてられないので記事化することにした．解説というよりは今やってることをそのまま書き散らしてるだけなので読みにくいかも．
環境の構成
GPUが乗ってるUbuntuマシン上でDockerを実行している． 研究のタスクごとにDockerのコンテナは切り分けていて， ひとつのタスクのために複数のコンテナを動かしてることもある．上図はひとつのタスクに2つのコンテナがあるイメージ．
各コンテナ上ではJupyterLabとTensorBoardが動いているので， コンテナ→ホストへフォワードしたポートにアクセスすれば作業を始めることができる．
マシン本体を直接使って作業することはほとんどなく， MacBook AirからSSHでマシンにアクセスしてリモートで作業を行う． SSHポートフォワーディングにより，コンテナ→ホスト→MacBookというようにポートをフォワードすることで，MacBookのブラウザからJupyterLabやTensorBoardを触ることができる．
SSH使わずにWebサーバとして公開するやりかたもあるけど，セキュリティ周りの管理とかだるくて今の形に落ち着いてる．
また，たまにインターンの学生さんに作業環境を引き継ぐことがあるので，そういうときにDocker使ってると対応が楽になって良い．
各コンテナ内のファイル構成
初期状態は以下のような感じ．
このホスト上に置かれたディレクトリは，
Dockerコンテナ内の/work以下にマウントされる．
このディレクトリのルートにJupyter Notebookをひたすら作っていく．
. ├── data ├── docker │ ├── docker-entrypoint.sh │ └── Dockerfile ├── docker-compose.yaml ├── log ├── modules └── README.md
| ファイル/ディレクトリ | 用途 |
|---|---|
| data | 使用するデータを突っ込んでおくディレクトリ．データの規模が大きいときは，このディレクトリの下に別のディレクトリをマウントすることがある． |
| docker/docker-entrypoint.sh | コンテナ起動時に実行したいコマンドを書いておくシェルスクリプト |
| docker/Dockerfile | 作業環境のDockerイメージを作るためのファイル |
| docker-compose.yaml | Dockerイメージのビルドやコンテナ起動時の設定を記述する． |
| log | TensorBoardによって監視させるディレクトリ |
| modules | 作ったり拾ってきたPythonモジュールを突っ込んでおくディレクトリ |
| README.md | なんのタスクをする環境なのかをメモしておく |
cookiecutterによる作業環境の自動生成
上記の作業環境を手動でコピペして増やすのは面倒くさいので， cookiecutterで簡単に環境を作れるようにした． (2019/9/19時点ではまだDockerfileにバグがあって，gidとuidが1000のユーザでないと使えない)
cookiecutterはPython製のコマンドラインツールで， 質問に答えていくだけで テンプレートからプロジェクトを爆速で生成できる．
機械学習の研究環境を生成するテンプレートは他にも cookiecutter-docker-scienceが知られているが， Anaconda環境で作業したかったり，docker-composeで 複数コンテナを同時に動かしたい事情があったので自前でテンプレートを 用意することにした．
ちなみにcookiecutter-docker-scienceの開発者はCookpadの方．
研究で作ったモデルやツールをPythonパッケージにする
開発した手法を誰かに使ってもらったり，自分で他のプロジェクトで使いまわしたりするときに pipでインストール可能な形式になっていると便利である．
自分はcookiecutter-pypackage-minimalでPythonパッケージを生成して， この上にモデルやデータの可視化のコードを書いてる．
生成したパッケージは/work/modules以下に突っ込んでおいて，
pip install -e <モジュール名>でインストールしておくことで，Jupyter Notebook側で動作確認しながら開発ができる．
研究の手法がきちんとパッケージの体裁になってると， 作業ディレクトリの設定が自由にできたり自分のプロジェクトに組み込みやすくて嬉しいのだけど， なかなか全部の論文でそうはなってない…．
綺麗にパッケージとして実装されてる深層学習の手法はopenpifpafなんかがある．
自分も研究とは関係ない個人的な開発で以下のpytorch-dataset-utilというパッケージをGitHub上で公開している． これは，特定のアノテーションツールで作成したアノテーションや torchvisionに実装されていないようなデータセットをPyTorchのDatasetクラスとして扱うインタフェースを提供している．
docker-composeによるコンテナの管理
Dockerのイメージをビルドしたりコンテナを起動/停止するにはdockerコマンドを使えば良いが，
複数のコンテナを同時に動かすようなプロジェクトだったりするとシェルスクリプトが肥大化して収拾がつかなくなってくる．
docker-composeはdocker-compose.yamlにDockerに関する必要な設定を書いておけばイメージやコンテナ名の設定，
コンテナの同時起動/停止/削除などをよしなにやってくれる便利コマンド．docker-composeもPython製のコマンドラインツールである．
自分の場合だと，前述したようなJupyterLabの作業環境に加えて，Webベースのアノテーションツールを同時起動することがある．
Gitでコードのバージョン管理・GitHubでコードのホスティング
開発したモデルやツールはみんなGitでバージョン管理している． Gitを使っておけば，「思い付きの方法を試してみていたらコードがだんだん散らかってきた」みたいな場面ですぐに前のバージョンに戻せる．
Jupyter Notebookまで管理するとわけがわからなくなるので，自分はNotebookは使い捨てにしてモジュールだけをGitで管理する使い方が多い．
GitHubは2019年から無料アカウントでもプライベートリポジトリを作成できるようになったので， 公開できないコードでも気軽に置けるようになった．
