MLflow 1.0.0 リリース！機械学習ライフサイクルを始めよう！

以前 Qiita で MLflow(ver0.4) に関する記事を書いたのですが、
最近(2019年5月22日)MLflow 1.0(候補版)がリリースされたらしいので再びまとめてみました。
本記事では MLflow の概要に加え MLflow1.0 + PyTorch を使ったコードを扱います。

MLflow 1.0 Released! | MLflow
Release MLflow 1.0.0 · mlflow/mlflow | github

0. 対象

• 機械学習アルゴリズム周りに関わっている人
• 頻繁に実験を回す人
• 機械学習を使ったちょっと長い期間のプロジェクトに所属してる人
• パラメータとか諸々の管理に疲れてきた人
• これから長期PoC案件にとりかかる人
• 混沌とした機械学習周りのなにかを引き継いでしまって悲しくなった人

1. MLflow とは

※ 図は公式サイトより引用

An open source platform for the machine learning lifecycle

MLflow は機械学習ライフサイクル(実験・再現・デプロイ)を支援するためのオープンソースプラットフォームです。以下の点をアピールしてます。

個人的には以下の点がポイント高いです。

• WORKS WITH ANY ML LIBRARY, LANGUAGE & EXISTING CODE
  • 機械学習フレームワークや言語に依存しない
    • API 機能周りは Java, R 等にも対応
  • 既存コードでも動く(影響が少ない)
• DESIGNED TO SCALE FROM 1 USER TO LARGE ORGS
  • 1人以上のメンバーにスケールできるように設計されている
    • 他の人と共有しやすいようになっている

事実多くのプラットフォームやフレームワークに対応しています(exampleがある)。

MLflow は大きく3つの機能を持っています。

• MLflow Tracking
  • 実験周りのコードや設定・結果の記録
• MLflow Projects
  • どこでも再現できるようにするためのパッケージング
• MLflow Models
  • モデルを各環境にデプロイするための方法やフォーマット

1.1 MLFlow Tracking

• コンセプト: 実験実行時の情報をまとめる
• 何を記録できるようになっているのか
  • Code Version(git の commit hash) → どの時点のどのコードなのか同定できるようになる
  • Start & End Time
  • Parameters(ハイパーパラメータとか)
  • Metrics(評価指標や epoch とか)
• どこに記録できるようになっているのか
  • ローカル
  • DB(specified as db_type://:@:/)
  • HTTP server
  • Databricks workspace(Azure のデータ分析プラットフォーム)
  • ストレージ
    • Amazon S3, Azure Blob Storage, Google Cloud Storage, FTP server
      , SFTP Serve, NFS
• 実験のグルーピング
  • タスクやプロジェクトを分別して同じ場所で管理できる
• Tracking UI(MLflow Tracking Servers)
  • サーバー上にプロセスを立ててアクセスできる(jupyter notebook みたいに)
  • local でもできる
  • UI を使ってまとめを確認できる
  • 結果のダウンロード等もできる

1.2 MLflow Projects

• コンセプト: 他の人などがデータ分析周りのコードが再利用・再現できるようにするためのパッケージング
  • YAML を使ったフォーマット
    • MLProject File というファイルを各実験用のリポジトリ内に用意する
  • プロジェクト名, スクリプト(実験コード)の実行方法, 実行環境

example-mlproject-file

name: My Project

conda_env: my_env.yaml
# Can have a docker_env instead of a conda_env, e.g.
# docker_env:
#    image:  mlflow-docker-example

entry_points:
  main:
    parameters:
      data_file: path
      regularization: {type: float, default: 0.1}
    command: "python train.py -r {regularization} {data_file}"
  validate:
    parameters:
      data_file: path
    command: "python validate.py {data_file}"

• Specifying Projects
  • 実行環境
    • conda, docker, 現在使っているシステム 等
  • 実行ディレクトリ
    • MLProject ファイルを含まないディレクトリ内を実行したい場合に細かく指定できる
  • MLProject File
    • 実行環境や各パラメータファイルの場所等をまとめた YAML 形式のファイル
    • 下記の例参照
• Running Projects
  • コマンドラインで mlflow run を使うか PythonAPI で mlflow.projects.run() で実行できる
  • 実行時に指定できる項目
    • Project URI
      • URL の広義版
      • https://<repo> (to use HTTPS) or user@host:path (to use Git over SSH)
    • Project Version
      • commit hash or branch name in the Git repository
    • Entry Point
      • 実行したいファイル(main.py や test.py, train.sh 等)を指定する
    • Parameters
      • ハイパーパラメータ等の実行時に与えるパラメータ
    • Environment
      • 基本は MLProject ファイル内の環境を使うが、他環境も指定できる
  • コマンド例(実際に動くらしいです)
    • mlflow run [e-mail]:mlflow/mlflow-example.git -P alpha=0.5
• Building Multistep Workflows
  • 「一回 main.py を回して終わり！」以外のケースにも対応している
    • 複数 git version を回したい
    • train, valid, test で step 分けたい
    • K-fold したい
    • ハイパーパラメータチューニングしたい

1.3 MLflow Models

• コンセプト: MLflow tracking で管理(保存)しているモデルを楽にデプロイするための保存フォーマット
  • デプロイ(提供)方法例
    • REST API としてリアルタイムに処理したい
    • Apache Spark 上でバッチ推論したい
  • ここでフレームワークによる保存方法などの違いも吸収する
• Storage Format
  • 各 MLflow Model はディレクトリ
    • MLmodel ファイルが root に置いてある
    • ディレクトリ構成例
      • my_model/
        • ├── MLmodel
        • └── model.pkl
  • MLmodel ファイルでは、モデルを参照するための各 flavor を定義できる
  • flavor とは(大事)
    • デプロイツールが保存したモデルの参照できるようにするためのルール
    • 各デプロイツールと各機械学習ライブラリの組み合わせごとに細かく書き直す必要がなくなる
    • モデルを Python 関数として実行する “Python function” flavor 等が標準で組み込まれている
    • 自分で定義することもできる
  • デプロイコマンド例(詳細コマンドは下記の PyTorch 例で触れます)
    • mlflow sklearn serve my_model
    • mlflow sagemaker deploy -m my_model [other options]
• Built-In Model Flavors(組み込みの flavor)
  • 一覧
    • Python Function (python_function)
    • H2O (h2o)
    • Keras (keras)
    • MLeap (mleap)
    • PyTorch (pytorch)
    • Scikit-learn (sklearn)
    • Spark MLlib (spark)
    • TensorFlow (tensorflow)
  • 上記のライブラリであればより低コストで MLmodel が導入できます
• Model Customization
  • 上記に含まれない機械学習ライブラリであっても、mlflow.pyfuncをベースにして組み込める
    • Python で使えるようにすれば MLflow に組み込める
  • ドキュメント内では xgboost を組み込んでました(参照リンク)
• Built-In Deployment Tools
  • 現在は全ての model flavors に対して全てのデプロイツールが対応していない
  • ただし Python Function フォーマットは全てのデプロイツールに対応している
  • ドキュメント内の例
    • Deploy a python_function model as a local REST API endpoint
    • Deploy a python_function model on Microsoft Azure ML
    • Deploy a python_function model on Amazon SageMaker
    • Export a python_function model as an Apache Spark UDF

2. 私たちにとって何が嬉しいのか

• 各実験を管理できる
  • あのときの結果は？
  • パラメータは？
  • 以前と何を変えたっけ？
  • あのときのコードは？
• 再現しやすい
  • あの環境でしか動かない・本番環境で再現できない、といったことを軽減できる
    • Dockerfile 等が環境を定義しているのに対し、MLflow では実行時のコマンドまで統一させようとしているイメージ
• API として提供しやすい
  • 実験・検証からデプロイまでの流れをスムーズにできる
  • 分担が別の人の場合にフォーマットですり合わせる手間が減る

3. 使ってみた(MLflow + PyTorch)

ここでは MLflow + PyTorch を使って MLflow tracking 周りを使ってみます。コードの詳細については触れないので、気になる方は githubリポジトリを参照ください。

注意1: 今回 MLFlow Tracking の機能のみ扱っています。
注意2: mlflow ui する場合は、mlflow リポジトリ外じゃないとうまく動きません(mlflow リポジトリを git clone してそのままサンプルコード動かそうとすると詰まる場合があります)

3.1 データセットやモデル

項目	値
データセット	CIFAR-10
モデル	ResNet18, SE-ResNext101
深層学習フレームワーク	PyTorch

3.3 事前準備(環境構築)

Mac の場合 brew install libomp しないと torch が動かない場合があります。

requirements.txt

torch==1.1.0
torchvision==0.3.0
pretrainedmodels==0.7.4
cnn-finetune==0.5.3
mlflow==1.0.0

Pythonパッケージインストール

$ python -V                                                                                                                                        
Python 3.7.3
$ pip install -r requirements.txt
...
Successfully installed mlflow-1.0.0

3.4 学習

学習時のおおまかな流れは以下のコードの通りです。
詳細な部分につきましては githubリポジトリをご参照ください。

with mlflow.start_run(): 以下の部分で各パラメータを保存しているのがポイントです。

main.py

def main():
    args = make_parse().parse_args()

    torch.backends.cudnn.benchmark = True
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    train_loader, test_loader = get_dataloaders(batch_size=args.batch_size,
                                                num_workers=args.num_workers)

    # num_classes and input_size is params of CIFAR-10
    print(f'model: {args.model}')
    model = make_model(args.model, num_classes=10,
                       pretrained=args.use_pretrain, input_size=(32, 32))

    criterion = torch.nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=args.init_lr, momentum=args.momentum)
    scheduler = optim.lr_scheduler.MultiStepLR(
        optimizer,
        milestones=[int(args.epochs*0.8), int(args.epochs*0.9)],
        gamma=0.1
    )

    if args.multi_gpu:
        model = torch.nn.DataParallel(model)
    model = model.to(device)

    results = train(
        epochs=args.epochs,
        model=model,
        train_loader=train_loader,
        valid_loader=test_loader,
        criterion=criterion,
        optimizer=optimizer,
        scheduler=scheduler,
        device=device
    )  # train_loss, val_loss, val_acc

    # mlflow
    with mlflow.start_run() as run:
        # Log args into mlflow
        for key, value in vars(args).items():
            mlflow.log_param(key, value)

        # Log results into mlflow
        for key, value in results.items():
            mlflow.log_metric(key, value)

        # Log other info
        mlflow.log_param('loss_type', 'CrossEntropy')

        # Log model
        mlflow.pytorch.log_model(model, "model")

学習時は以下のようにパラメータを少し変えて2通り試してみます。
--multi-gpu は GPU 1枚のときは外してください。
また se_resnext50_32x4d の pre-trained モデルはダウンロード遅いので、気になる人は外してください。

train.sh

python main.py \
    --batch-size 1024 \
    --num-workers 16 \
    --init-lr 1e-2 \
    --momentum 0.9 \
    --epochs 30 \
    --model resnet18 \
    --use-pretrain \
    --multi-gpu

python main.py \
    --batch-size 1024 \
    --num-workers 16 \
    --init-lr 1e-2 \
    --momentum 0.9 \
    --epochs 30 \
    --model se_resnext50_32x4d \
    --use-pretrain \
    --multi-gpu

学習実行

$ sh train.sh
...
Epoch 30 loss = 4.66466e-07 val_acc = 0.7848

mlruns ファイルがあれば OK です。

3.6 結果確認

$ mlflow server -h 0.0.0.0 -p 8888
[2019-06-05 13:30:05 +0000] [24084] [INFO] Starting gunicorn 19.9.0
[2019-06-05 13:30:05 +0000] [24084] [INFO] Listening at: http://0.0.0.0:8888 (24084)
[2019-06-05 13:30:05 +0000] [24084] [INFO] Using worker: sync
...

local で実行中の場合は localhost:8888 にアクセスしましょう。
下記のような画面が出てくれば OK です！
(下図のアドレスが違うのは、docker コンテナ上で実行しているためです)

下図のように見たい項目だけ入力して Search すると絞り込めるので便利です！

4. 使ってみた感想

• パラメータ周りを固定フォーマットで出力するようになるので長期で運用するとメリットが大きそう
  • まめな人は既にここら辺管理していると思うのですが、他人と共有する前提で管理してる人は少数派だと思います
• 本家の example が色々あって参考になりました
• git リポジトリのバージョンと対応させられるのは便利そう
• まずは個人で導入しやすいところから始めるのが良いと思います
  • 個人 → プロジェクト → 部署 等
  • 人数や規模に対してスケールしやすそうですし
• 個人で書き捨てる予定なら無理して使う必要はないと思います
  • (とはいえ明日や昨日の自分は他人と思えば・・・)
• 始めたい場合は MLFlow Tracking を個人で始めるのが良さそう
• 評価指標が途中で変わった場合にどのようなフローで進めるのかよぐわがらない

最後に

MLflow は便利そうだしいろんな機能がある一方、運用するためには様々なルールに従う必要があります。ですので今まで好き勝手してきた人(自分とか)は「なにこれめんどくさ！」と思うかもしれません。

ただし今日のデータ分析は１人で全てを完結するのは難しく、他の人(過去や未来の自分含め)といかに共有して進められるかがより大事になってくると思います。

必ずしも MLflow を使うべきとは思いませんが、MLflow に触れてみることで機械学習ライフサイクルを回す上でどのような要素が必要・不要なのかを考えてみることは役に立つと思います。

皆さんステキな機械学習ライフサイクルを！

References

• MLflow - A platform for the machine learning lifecycle
  • 本家サイト
• mlflow/mlflow: Open source platform for the machine learning lifecycle
  • 本家コード
  • まだ chainer 周りが整理されてないので皆さん OSS contribute のチャンスです！
• mlflowを使ってデータ分析サイクルの効率化する方法を考える - Qiita
  • 分かりやすい記事です
• MLflow 〜これで機械学習のモデル管理から API 作成まで楽にできるかも〜 - Qiita
  • 1年前の自身の記事です
  • ver 0.4 なので注意