aimlflowによるMLflowとAimの連携

はじめに

会社の上司からaimlflowというものがあるよと紹介されました。実験管理ではMLflowを良く使っているのですが、このUIを改善してくれそうだというので、ちょっと調べてみました。

aimlflowとは

端的に言うと、MLflow Trackingの機能で記録された実験結果をAimというツール上で表示できる形式に変換してくれるのがaimlflowの役割になります。

AimはMLにおける実験管理機能を提供するOSSです。その点ではMLflow Trackingと機能は被っています。AimはMLflowに比べるとUIの表現力が高いと感じます。一方で、サポートするフレームワークの多さやModel Registryなどのコンポーネントも有している点で、MLflowは機能性が高いと感じます。

aimlflowは、実験管理としてMLflowを利用しているがその結果をよりビジュアルに確認したい、というケースで利用できると考えています。

操作

1. MLflowの準備

まずはpython仮想環境を作ってMLflowをインストールして、Trackingサーバーが稼働することを確認します。

事前準備

# python仮想環境の作成と有効化
python -m venv venv
source venv/bin/activate

# MLflow準備
mkdir mlruns
touch mlruns/mlruns.db
pip install mlflow

# MLflow Trackingサーバーの起動
mlflow server --host 0.0.0.0 --port 5000 \
    --backend-store-uri sqlite:///mlruns/mlruns.db \
    --default-artifact-root mlruns

2. 学習プログラムの作成

MLflowにトラッキングさせる学習プログラムsrc/train.pyを作成します。

src/train.py

import sys
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import ElasticNet
import mlflow

if __name__ == "__main__":
    # データのダウンロード
    csv_url = "http://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
    data = pd.read_csv(csv_url, sep=";")

    # MLflow設定: Experimentの指定とAutologの設定
    mlflow.set_tracking_uri("http://localhost:5000")
    experiment = mlflow.set_experiment("experiment_winequality-red")
    mlflow.sklearn.autolog(registered_model_name="winequality-red-model")

    # データをトレーニングセットとテストセットに分割
    train, test = train_test_split(data)
    train_x = train.drop(["quality"], axis=1)
    test_x = test.drop(["quality"], axis=1)
    train_y = train[["quality"]]
    test_y = test[["quality"]]

    # alphaとl1_ratioの値を定義(値はコマンドライン引数で受け取る：デフォルト値は共に0.5)
    alpha = float(sys.argv[1]) if len(sys.argv) > 1 else 0.5
    l1_ratio = float(sys.argv[2]) if len(sys.argv) > 2 else 0.5

    # 学習を実行
    lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
    lr.fit(train_x, train_y)

プログラムが用意できたら、引数を変えて何回か実行してみます。

train.pyの実行

python src/train.py 0.5 0.5
python src/train.py 0.4 0.4
python src/train.py 0.3 0.3
python src/train.py 0.2 0.2
python src/train.py 0.1 0.1

上記のように実行すると、5回分の実行結果がMLflowに記録されています。MLflow UI(http://localhost:5000)で実行結果を確認できます。

3. Aimの準備とaimlflowの実行

aimlflowをインストールします。インストールはpipで行えます。また、aimlflowと同時にaimも一緒にインストールされます。

aimlflowのインストール

pip install aim-mlflow

次に、Aimを初期化します。下記コマンドを実行することで、.aimディレクトリが生成されます。このディレクトリがAimの実験管理のリポジトリとなります。

Aim初期化

aim init

ここまできたら、aimlflowを実行します。この際、MLflow Trackingサーバーは稼働したままにしておき、下記コマンドを実行します。--mlflow-tracking-uriにはMLflow TrackingサーバーのURL、--aim-repoには.aimディレクトリが存在するパスを指定します。

aimlflowの実行

aimlflow sync --mlflow-tracking-uri=http://localhost:5000 --aim-repo=.

コマンド実行すると下記のようなメッセージがターミナルに表示されます。

実行結果

・・・

andler for html file types is not yet implemented.equality-red`:  60%|██████████████████████████████                    | 3/5 [00:00<00:00,  6.75it/s]
Unresolved or unsupported type for artifact model/MLmodel
Unresolved or unsupported type for artifact model/model.pkl

andler for html file types is not yet implemented.equality-red`:  80%|████████████████████████████████████████          | 4/5 [00:00<00:00,  6.87it/s]
Unresolved or unsupported type for artifact model/MLmodel
Unresolved or unsupported type for artifact model/model.pkl
Parsing mlflow runs for experiment `experiment_winequality-red`: 100%|██████████████████████████████████████████████████| 5/5 [00:00<00:00,  6.81it/s]
Parsing mlflow runs for experiment `Default`: 0it [00:00, ?it/s]█████████████████████████                               | 1/2 [00:00<00:00,  1.28it/s]
Parsing mlflow experiments in http://localhost:5000: 100%|██████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  2.53it/s]
Starting watcher on http://localhost:5000.

4. Aimでの実験結果確認

下記コマンドを実行してAimを起動します。

Aim起動

aim up

起動したら、aimのUIにアクセスします。URLのデフォルトは http://localhost:43800 になります。下図のような画面が表示され、MLflowに記録されていた実験結果がAimのUIで確認できることが分かります。

おわりに

今回はMLflowの結果をAimで見れるようにしました。あとはAimでどのような結果の見せ方ができるのかになりますが、これはまた別の機会に整理したいと思います。