Sansan Tech Blog

Sansanのものづくりを支えるメンバーの技術やデザイン、プロダクトマネジメントの情報を発信

【Intern CV Report】Trainsによる実験管理

R&D インターン 連載

f:id:S_aiueo321:20190404193803p:plain

こんにちは，DSOC 研究開発部 インターン生の内田です．先日プチ卒業旅行として伊豆の温泉旅館に行ってきました．伊豆といえば バナナワニ園 ですよね，異論は認めません．バナナワニ園のワニたちは基本的に動かないので思わず「休日の俺じゃん」って呟いてしまいました．

さて，今回はCVとは直接関係ないですが，多くの人が頭を悩ませている実験管理に関する話題です．

機械学習と実験管理

機械学習，特に深層学習を用いたプロジェクトにおいては，精度を引き出すためにハイパーパラメータが非常に重要です．一口にハイパーパラメータといっても，モデルサイズや学習率，バッチサイズなど挙げるとキリがありません．

最適なハイパーパラメータを定めるには，Grid Searchなりでパラメータを探索し，精度などと共にロギングする必要があります．バナナワニ園のワニのように怠惰な私にとって，パラメータや結果をマメに記録しておくことは簡単なことではありません．現状では表示可能な項目はなるべくTensorBoardに集約して，TensorBoardのログディレクトリにその他の結果を全部吐かせるようにしています．これなら後からでも漏れなく参照できるので安心ですが，TensorBoardに表示できない項目の検索・比較は面倒です．

上記の背景から，もう少し高機能な実験管理ツールを使ってみたい気持ちが常々あります．ただし，導入にあたってはいくつか条件があります．パッと思い当たるのは次のような感じです．

  • 移行コストが低いこと
    ロガーに値を手動で報告する場合，ツールを変えれば当然プログラムの書き換えが生じます．一気に書き換えるのも手ですが，バックアップとしてローカルのログも残しておきたかったりもするので，ロガーの命令をキャッチしてよしなにやってくれると嬉しいです．

  • モデル管理機能
    学習済みモデルと実験を紐づけてUI上で管理できると嬉しいです．

  • 社内サーバでのホスティング
    最近 Comet.ml や Neptune.ai のようなクラウド上で実験管理を行えるSaaSが出てきていているのですが，業務で扱うデータの性質上，外部のサーバにデータを送信するのはご法度です．とはいえ個々人のインスタンスに閉じているのもスケールしなさそうなので，社内サーバにデプロイして共有できる程度が望ましいです．

Trains

上記の条件に概ね合致しているツールとしてTrainsがあります．Trainsは，Allegro AI社が提供する実験管理ツールであり，OSSとしてGitHubに公開されています．

github.com

Trainsそのものはコード内から結果・モデルを報告するPythonパッケージであって，pip install trainsでインストールできます．学習曲線やハイパーパラメータはもちろんのこと，モデルのスナップショットやコミットIDなども自動でロギングしてくれます．2行プログラムを追加するだけでTensorBoardの命令をキャッチして自動でロギングしてくれるので，移行コストが非常に低いのも特徴です．

また，Trainsから報告された結果は，実験管理UIを提供するTrains Server上でホスティングされます．社内サーバにTrains Serverをデプロイしておけば，複数プロジェクトを1つのUI上で管理することが可能になります．

半年くらい前にTwitter上でちょっと話題になっていて気になってはいたのですが，あまり使うタイミングがないままだったので，今回はドキュメントを読みながらTrains Serverの導入からTrainsの使い方についてまとめたいと思います．

実行環境

本エントリでは以下の環境を想定しています．

  • OS: Ubuntu 16.04 LTS
  • Docker: 19.03.0
  • Docker Compose: 1.24.1
  • Python: 3.7.4
  • Trains / Trains Server: 0.14

Trains Serverの導入

Trains Serverはビルド済みイメージが配布されているため簡単に導入できます．以下ではDockerを利用しますが，他にもEC2 AMIが利用できたり，k8sクラスタ上にも導入できます．Trains Serverのシステム概略図は次のようになります．直近のリリースではこれに加えてredisのコンテナも動いているようで，日々改善されている感があります．

f:id:S_aiueo321:20200316022454p:plain
Trainsのシステム概略図*1

導入手順は以下です．

  1. TRAINS ServerではElasticSearchを利用するため，メモリマップ数の上限を引き上げます．

    echo "vm.max_map_count=262144" > /tmp/99-trains.conf
    sudo mv /tmp/99-trains.conf /etc/sysctl.d/99-trains.conf
    sudo sysctl -w vm.max_map_count=262144
    sudo service docker restart

  2. データベース・ストレージ用のディレクトリを作成します．

    sudo mkdir -p /opt/trains/data/elastic
    sudo mkdir -p /opt/trains/data/mongo/db
    sudo mkdir -p /opt/trains/data/mongo/configdb
    sudo mkdir -p /opt/trains/data/redis
    sudo mkdir -p /opt/trains/logs
    sudo mkdir -p /opt/trains/config
    sudo mkdir -p /opt/trains/data/fileserver

  3. 作成したディレクトリ群へのアクセス権をDockerに付与します．

    sudo chown -R 1000:1000 /opt/trains

  4. GitHubから docker-compose.yml をダウンロードします．

    cd /opt/trains
    curl https://raw.githubusercontent.com/allegroai/trains-server/0.14.0/docker-compose.yml -o docker-compose.yml

  5. Docker Compose でサーバを起動します．

    sudo docker-compose -f docker-compose.yml up

以上でTrains Serverのデプロイは完了で，下記のポートが利用可能となります．

  • 8080: Webサーバ
  • 8008: APIサーバ
  • 8081: ファイルサーバ

localhost:8080 にアクセスし，次のようなダッシュボードが表示されれば成功です．

f:id:S_aiueo321:20200316033145p:plain
Trains Web-Appのダッシュボード

Trains Serverへの送信設定

PythonコードからTrains Serverへデータを送信するには，認証情報の登録を行う必要があります *2． ダッシュボード左下のメニューからPROFILEを選択し，Create new credentials をクリックすると次のような画面になります．

f:id:S_aiueo321:20200316022810p:plain
認証情報の取得

アクセスキーとシークレットキーが生成されているので，Copy to clipboard をクリックして認証情報を取得します．コンソールからtrains-initコマンドを実行すると認証情報を聞かれるので，クリップボードの情報を貼り付けて設定完了です．設定項目は~/trains.confに保存されているので，二回目以降は~/trains.confを直に編集すれば設定を変更できます．

Trainsによる実験管理

以下では，PyTorch Lightningを用いてニューラルネットワークを学習し，Trainsによって実験管理を行います．PyTorch LightningはデフォルトでTensorBoardによるロギングを行うため，最小限のコード追加でTrainsによる自動ロギングが可能です．

実験では次のコードを用います．

import argparse
import os

import pytorch_lightning as pl

from models import CoolModel

# Trains関係の追加部分は次の2行のみ
import trains
task = trains.Task.init(project_name="trains-demo",　task_name="mnist")


def app():
    parser = argparse.ArgumentParser()
    parser.add_argument('--n_filters', type=int, required=True)
    args = parser.parse_args()

    trainer = pl.Trainer(max_epochs=10)

    model = CoolModel(args)

    trainer.fit(model)

if __name__ == "__main__":
    app()

CoolModelは3層ニューラルネットワークによりMNISTを学習するモデルクラスであり，コマンドラインから隠れ層のフィルタ数をハイパーパラメータとして受け取ります．Trainsに関連する部分はコメント以下2行のみです．上記のコードを5種類のフィルタ数に対して実行すると，ダッシュボードは次のように変化します．trains.Task.init()で指定したように，trains-demoというプロジェクトが出現し，プロジェクト下にmnistという実験が5つ生成されているのが確認できます．

f:id:S_aiueo321:20200316022817p:plain
実験後のダッシュボード

各実験の詳細は，プロジェクトに入りポチポチすることで確認できます．各実験はユニークなIDで管理されており，IDボタンをクリックすると取得できます．

f:id:S_aiueo321:20200316022823p:plain
実験の詳細画面

各タブには次のような項目が記録されています．

大項目 小項目 内容
EXECUTION - コミットID，依存パッケージ等の再現性に関する情報
HYPER PARAMETERS - ハイパーパラメータ
ARTIFACTS - モデルファイルなどの入出力ファイル
INFO - 実験時間等のメタデータ
RESULTS LOG 標準入出力
RESULTS SCALARS 学習曲線やマシンリソースモニタリング
RESULTS PLOTS MatplotlibやSeabornで作成した図
RESULTS DEBUG IMAGES 画像ロギング

また，各実験のチェックボックスを複数選択することで，実験間でハイパーパラメータや学習曲線を比較できます．ベースの実験との差分をハイライトしてくれるため，比較が容易です．

f:id:S_aiueo321:20200316022830p:plain
実験間のハイパーパラメータ比較画面

f:id:S_aiueo321:20200316022758p:plain
学習曲線の比較．曲線の色や表示数などお好みで制御可能．

Trainsによるモデル管理

Trainsはモデル管理機能も提供しています．実験管理と同様に，モデルファイルの保存を自動的にキャッチして記録を取ってくれます．例えば次のコード*3を実行すると，プロジェクトのモデルタブに出力モデルが記録されます． ここで Task.init(..., output_uri="s3://<BUCKET>") みたいに指定しておくと，クラウドストレージにモデルを転送することもできます *4．

import torch
from torchvision.models import alexnet

from trains import Task
task = Task.init(project_name='trains-demo', task_name='model_saving')

model = alexnet()
torch.save(model.state_dict(), 'model.pth')

モデル管理画面は次のようになります．実験とモデルの紐付けはCREATING EXPERIMENTから確認できます．

f:id:S_aiueo321:20200316025635p:plain
モデル管理画面

実験と同様，モデルはユニークIDで管理されるため，IDボタンをクリックしてモデルIDを取得できます． モデルIDから保存先のパスを取得するには，次のように記述します． 保存先がクラウドストレージである場合，バケットからモデルファイルがtempディレクトリにダウンロードされ，ダウンロード先のパスが返ってきます．

from trains import InputModel
model_path = InputModel('model_id').get_weights()

ファイルの保存先にかかわらずモデルIDのみで読み出せるので，コードの簡素化も期待できますね．

まとめ

実験管理ツール Trains を紹介し，実験管理からモデル管理までの流れを実際に試してみました． UIが非常に良くできていますし，自動ロギングやクラウドストレージとの連携がスムーズだったので実用性があると感じました． 特に，今までテストの際にモデルファイルのパスを手打ちしていたことを考えると，モデルIDをUIから取得すればいいのはとても嬉しく感じました．

また今回紹介できませんでしたが，「既存の実験をクローン→ハイパーパラメータを修正→再実行」というフローを実現するTrains Agentというツールも存在します． 発展著しいツールですので，普及に期待したいと思います．


▼本連載のほかの記事 buildersbox.corp-sansan.com

*1:https://allegro.ai/docs/concepts_arch/concepts_arch/ より引用

*2:ここで設定を行わない場合，データはデモサーバに送られ，誰でも見れる状態に置かれます．

*3:Lightningのモデル保存方法が特殊でうまくキャッチしてくれなかったため変更しています．

*4:アクセスキーを trains.conf に設定しておく必要があります．

© Sansan, Inc.