MLflow: 1. Tracking を使った実験管理

はじめに
- 環境情報
MLflowとは
MLflowでできること
MLflow Trackingとは
- トラッキングできる情報
- 概念の説明
  - エクスペリメントとは:
  - ランとは:
基本的な操作
まとめ
参考

はじめに

多くの機械学習プロジェクトでは、精度の向上のためパラメータや特徴量を変えて試行錯誤しながら、何度も実験を繰り返す必要がある。そのため試行回数が増加するごとにモデルの管理が難しくなり、属人化するという課題がある。そこで今回はMLflow Tracking (トラッキング) を使用したモデルの管理についてまとめる。

環境情報

Databricks RunTime: 10.2 ML (includes Apache Spark 3.2.0, Scala 2.12)

MLflowとは

機械学習のライフサイクルをエンドツーエンドで管理するためのオープンソースのプラットフォーム。データの準備からモデル開発、デプロイ、モデルの監視までの機械学習モデルのライフサイクルを一貫して管理することを目標としている。2022年02月時点でGithub上で11.4Kのスターを獲得しており、機械学習のモデル管理ツールとして広く認知されているツールである。

MLflowでできること

MLflowは4つの主要な機能から構成されており、各機能は独立して使用することが可能である。

1. MLflow Tracking:

実験をトラッキングし、パラメータや結果を記録・比較する。

2. MLflow Projects:

他のデータサイエンティストとの共有や本番環境への移行のために、MLコードを再利用可能で再現性のある形でパッケージ化する。

3. MLflow Models:

多様なデプロイツールをサポートする汎用的なモデル形式を提供する。

4. MLflow Model Registry:

モデルのバージョン管理、ステージ遷移、アノテーションなど、MLflowモデルの全ライフサイクルをセントラルモデルストアとして一貫して管理する。

今回はMLflowの最も基本的な機能であるMLflow Trackingについて詳しく説明する。

MLflow Trackingとは

機械学習コードの実行時にパラメータ、コードバージョン、メトリクス、出力ファイルをロギングし、結果を可視化するためのAPIとUIのこと。MLflow Trackingでは、Python、REST、R API、Java APIの各APIを使って、実験のログやクエリを行うことができる。 MLflow Trackingは、機械学習コードの1度の実行であるランという概念を中心に構成されており、各ランには以下の情報が記録される。

トラッキングできる情報

・実行時間: 一度のランを実行するのにかかった時間。
・実行者: ランを実行したユーザ名。
・ソース: トラッキングしたランが含まれるノートブックの名前。MLflow Projectから実行する場合は、プロジェクト名と実行のエントリーポイント。
・バージョン: ノートブックから実行した場合は、ランの実行時のノートブックバージョン(履歴)。MLflow Projectから実行する場合は、実行に使用されたGitコミットのハッシュ値。
・パラメーター: キーと値のペアとして保存されるモデルのパラメーター。（キーと値はどちらも文字列であることに注意する）
・メトリクス: キーと値のペアとして保存されるモデル評価のためのメトリクス。（値は数値のみ記録可能であることに注意する）各メトリクスは、例えばモデルの損失関数の収束のトラッキングなど各実行の過程を記録・更新することができる。そのためMLflowはメトリクスの全履歴を記録し、視覚化することが可能である。
・タグ: キーと値のペアとして保存されるメタデータ。（キーと値はどちらも文字列であることに注意する）
・アーティファクト: 学習済みモデルや、特徴量の重要度など実験した結果として得られる成果物。任意のフォーマットでファイルの出力ができる。例えば、画像（PNGなど）、モデル（pickleしたscikit-learnモデルなど）、データファイル（Parquetファイルなど）を成果物として記録することが可能である。

概念の説明

エクスペリメントとは:

各ランをノートブックまたはワークスペース単位でまとめたもの。ノートブック単位でまとめた場合は、同じノートブック内の全てのランが1つのエクスペリメントに紐付く。一方で、ワークスペース単位でまとめた場合はノートブックに関連付けられるのではなく、どのノートブックのランでも実験IDまたは実験名を使用することで、同じ1つのエクスペリメントに紐づけることが可能。

ランとは:

各モデルの 1 回の実行に対応した試行のこと。例えばxgboostで任意のデータセットを学習させる試行を1度行った場合は、それが1つのランに対応する。

基本的な操作

各ランをトラッキングする方法は以下の2種類がある。自分でロギング内容を定義するのは面倒なので、基本はautolog()関数を使った自動ロギングがおすすめ。

ロギング関数を使って必要な情報をロギングする
自動ロギングを使って必要な情報を自動でロギングする

1. ロギング関数を使って必要な情報をロギングする

以下の関数を使って、各ランで必要な情報をロギングする。
※関数の詳細はMLflow Tracking — MLflow 2.9.2 documentationを参照。

MLflow のパッケージをインポートする。

import mlflow

ロギングを開始する。

with mlflow.start_run(run_name='hogehoge_model'):

パラメータを記録する。

mlflow.log_param('n_estimators', n_estimators)

メトリックとしてaucを登録する。

 mlflow.log_metric('auc', auc_score)

タグを登録する。

mlflow.set_tag(key='hoge', value='fuga')

アーティファクトを登録する。

import mlflow

# Create a features.txt artifact file
features = "rooms, zipcode, median_price, school_rating, transport"
with open("features.txt", 'w') as f:
    f.write(features)

# With artifact_path=None write features.txt under
# root artifact_uri/artifacts directory
with mlflow.start_run():
    mlflow.log_artifact("features.txt")