aimlflowによるMLflowとAimの連携


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

https://github.com/aimhubio/aimlflow

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

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

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

# 操作
## 1. MLflowの準備
まずはpython仮想環境を作ってMLflowをインストールして、Trackingサーバーが稼働することを確認します。
```sh:事前準備
# 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`を作成します。
```python: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)
```

プログラムが用意できたら、引数を変えて何回か実行してみます。
```sh: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)で実行結果を確認できます。

![mlflow-1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2724847/6a1b960d-3253-19fe-568c-401f9a14fd10.png)


## 3. Aimの準備とaimlflowの実行
aimlflowをインストールします。インストールはpipで行えます。また、aimlflowと同時にaimも一緒にインストールされます。
```sh:aimlflowのインストール
pip install aim-mlflow
```

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

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

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

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を起動します。
```sh:Aim起動
aim up
```

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


![aim-1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2724847/62be3a0d-6c4d-0697-56ff-d6720a1fa198.png)


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