【実践編】MLflow Trackingによるモデル実験管理

こんにちは、ガロガロニーです。本日は機械学習モデル開発の実験管理について、MLflow Trackingを利用した例を紹介します。

機械学習における再現性の課題
MLflowとは
環境構築
モデル実験管理の例
MLflowコード
まとめ

機械学習における再現性の課題

機械学習のモデル開発では、複数人が様々な条件に対して、何度も試行を繰り返しながら実験を行います。そのため、実験を進めていくうちに徐々に管理が難しくなってしまい、後から結果を再現したくてもできない問題が発生します。機械学習において再現性というのはとても重要です。せっかく実験で良い結果が得られても、あとから結果が再現できなければ意味がありません。結果が再現できなかれば、そのモデルの確からしさを損なうことになります。実際の現場でも、モデルの実験管理ができておらず、クライアントに説明したモデルの精度が再現できなかったり、本格的なシステム化・製品化の検討ができなかったりするケースが非常に多くなっています。

MLflowとは

MLflowは、Databricks社がオープンソースで提供している機械学習パイプラインのOSSで、ローカル環境で軽量に利用することができるプラットフォームです。Apache License 2.0でgithub上で公開されています。MLflowには4つの主要な機能があります。

MLflow Tracking：実験関連のコード、データ、設定や結果の記録
MLflow Projects：データサイエンスコードのパッケージ化
MLflow Models：モデルを各環境にデプロイ、API作成できる機能
MLflow Registry：モデルを保存、管理できるレジストリ

今回はMLflowの機能の中でも、モデルの実験管理を簡単に行うことができるMLflow Trackingについて、実際の使用例を紹介していきます。

環境構築

MLflowはpipがインストールされていれば、pip install mlflowで簡単に導入できます。また機械学習の実験を行うため、sklearn、matplotlib、xgboost、lightgbmもインストールしておきます。今回のPythonやライブラリのバージョンは下記の通りです。（2022年1月1日時点での最新バージョンを使用しています。）

python：3.9.7
mlflow：1.22.0
sklearn：1.0.2
matplotlib：3.5.1
xgboost：1.5.1
lightgbm：3.3.1

次に、github からMLflowのzipファイルをダウンロードし、任意のディレクトリに展開してください。今回はMLflowの公式コードにあるexamplesフォルダの内容を使っていきます。

githubからMLflowコードをダウンロードし、展開した時のディレクトリ構成

モデル実験管理の例

様々な機械学習アルゴリズムを使って、モデル開発を行った際の実験管理の例を紹介します。今回はXGBoostやLightGBMを使ってみたいと思います。これらは勾配ブースティングフレームワークと呼ばれるもので、様々なデータに対する分類や回帰の分析で使われます。機械学習のコンペでも非常に人気でよく採用されています。MLflowはXGboostやLightGBMなどの有名なフレームワークに対応しており、簡単に実験管理を行うことができます。今回はこれらのフレームワークを使い、irisと呼ばれるアヤメの種類と特徴量に関するデータセットを使って、アルゴリズムやパラメータを変えながらモデルの構築を行い、MLflowで実験管理できる例を紹介します。

XGBoost

まずはXGBoostを使って、モデル構築を行います。学習パラメータを変えながら、実験を2回行います。

XGBoostの学習パラメータをコマンドライン引数で渡し、モデルの実験（1回目）を行います。モデル構築します。
python examples\xgboost\xgboost_native\train.py --learning-rate 0.2 --colsample-bytree 0.8 --subsample 0.9
続いて、学習パラメータを変更して、モデルの実験（2回目）を行います。
python examples\xgboost\xgboost_native\train.py --learning-rate 0.4 --colsample-bytree 0.7 --subsample 0.8

LightGBM

続いてLightGBMにアルゴリズムを変えて、モデル構築を行います。同様に学習パラメータを変えながら、3回目と4回目の実験を行います。

LightGBMの学習パラメータをコマンドライン引数で渡し、モデルの実験（3回目）をしてみます。
python examples\lightgbm\train.py --learning-rate 0.2 --colsample-bytree 0.8 --subsample 0.9
続いて、学習パラメータを変更して、モデルの実験（4回目）を行います。
python examples\lightgbm\train.py --learning-rate 0.4 --colsample-bytree 0.7 --subsample 0.8

MLflow UIによる実験の可視化

XGBoostとLightGBMを使ったモデル構築の実験結果を可視化します。mlflow uiというコマンドを実行し、http://127.0.0.1:5000をブラウザで開いてみてください。実験を行った時間やアルゴリズム、コード、評価結果、パラメータなどが実験ごとに記録されていることが分かります。

さらに、個別の実験結果の詳細についても確認することができます。学習済みモデルや実行環境、モデルの特徴量重要度（Feature Importance）などの情報を把握することができます。

MLflowコード

MLflowによる実験管理は、コード上でも簡単に記述することができます。下記はXGBoostの実験で使用したコードの一部ですが、mlflow.xgboost.autologメソッドを使えば自動でロギングを行ってくれます。LightGBMについても同様に、mlflow.lightgbm.autologメソッドを使うことができます。自動ロギングは下記の情報を自動で記録してくれます。

trainメソッドで指定した学習パラメータ
evalsで指定した誤差推移
学習済みモデルファイル
特徴量重要度（Feature Importance）

なお評価結果に関しては、指標がタスクごとに異なるため、手動で設定する必要がありますが、mlflow.log_metricsメソッドを使用すれば、簡単に記録することができます。

まとめ

本日はMLflow Trackingを利用したモデル実験管理について、XGBoostやLightGBMによる実例を交えながら紹介しました。MLflowは軽量なプラットフォームで簡単に導入でき、これまで苦労していたモデル開発の再現性の課題を改善することができます。特にチームで開発していく機械学習プロジェクトでは、実験管理がとても重要になってくるため、ぜひ今回紹介したMLflowを試してもらえると嬉しいです。それではまた！