Databricksにおける機械学習モデルとトレーニングデータのトラッキング

以前、こちらの記事を書きました。

https://qiita.com/taka_yayoi/items/523c0b1b52e3b39292b1

機械学習モデルトレーニングの管理においては、**トレーニングでどのデータを使ったのか**は再現性確保の観点で重要です。しかし、手動でそのような情報を管理するのにも限界があります。

上述の記事では、MLflowでモデルをトラッキングする際に、トレーニングに用いたテーブルをバージョンとともに記録できる機能を説明しました。1年以上前の記事ということもあるので、機能を再度確認します。

こちらでは、以下の観点でウォークスルーします。

1. MLflowエクスペリメントにおけるモデルとデータのトラッキング
1. Unity Catalogにおけるモデルとデータのトラッキング

[MLflowのエクスペリメント(実験)](https://learn.microsoft.com/ja-jp/azure/databricks/mlflow/experiments)とはMLflowの用語ですが、MLflowにおいて機械学習モデルトレーニングをトラッキングする際、個々のトレーニングは[MLflowラン(実行)](https://learn.microsoft.com/ja-jp/azure/databricks/mlflow/runs)として管理します。エクスペリメント複数のMLflowランをまとめて管理するための箱のようなものです。これらは、名前の通り、実験を試行錯誤している段階でモデルに関連する情報を記録するためのものです。

そして、Unity Catalogでのトラッキングに言及しているのは、実験から本格運用に移行する際に[Unity Catalogでモデルのライフサイクルを管理](https://learn.microsoft.com/ja-jp/azure/databricks/machine-learning/manage-model-lifecycle/)することになるためです。以前は、この機能は[MLflowのモデルレジストリ(ワークスペースモデルレジストリ)](https://learn.microsoft.com/ja-jp/azure/databricks/machine-learning/manage-model-lifecycle/workspace-model-registry)が担っていましたが、DatabricksではUnity Catalogのモデルレジストリに相当する機能(Models in UC)を使うことを推奨するようになっています。

# MLflowエクスペリメントにおけるモデルとデータのトラッキング

最新のランタイムを使えば、[前回の記事](https://qiita.com/taka_yayoi/items/523c0b1b52e3b39292b1)で言及した機能を活用するためのMLflowは最初から利用できます。
![Screenshot 2024-07-17 at 8.35.22.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/bf59e894-3370-06cb-68fb-32bda0fb65fb.png)

## Auto Loggingによるトレーニングデータのトラッキング

[こちらのサンプル](https://learn.microsoft.com/ja-jp/azure/databricks/machine-learning/manage-model-lifecycle/#register-a-model-to-unity-catalog-using-autologging)を一部変更しながら機能を確認します。

```py
from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier

# irisデータセットでsklearnモデルをトレーニング
X, y = datasets.load_iris(return_X_y=True, as_frame=True)
clf = RandomForestClassifier(max_depth=7)
clf.fit(X, y)
```

[Auto Logging](https://learn.microsoft.com/ja-jp/azure/databricks/mlflow/databricks-autologging)によってscikit-learnのモデルなどは自動でロギングされます。
![Screenshot 2024-07-17 at 8.38.10.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/1adbd895-4570-a227-348b-ad02bec12049.png)

エクスペリメント配下のラン詳細を確認します。
![Screenshot 2024-07-17 at 8.38.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ee3f0906-38c2-3672-96bc-4b7d5071a6fd.png)

この時点で**使用されたデータセット**からトレーニングデータを確認することができます。
![Screenshot 2024-07-17 at 8.39.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/bac2c000-9a41-3e6a-dc55-6f0d6dc7b29a.png)

## `mlflow.log_input`によるトレーニングデータテーブルのトラッキング

上の例でも、トレーニングデータ自体はトラッキングできていますが、より厳密にトレーニングデータを管理するのであれば、Unity Catalogのテーブルでトレーニングデータを保持すべきです。耐障害性、バージョン管理、アクセスコントロールなどさまざまな面でのメリットを享受することができます。

[`mlflow.log_input`](https://mlflow.org/docs/latest/python_api/mlflow.html?highlight=log_input#mlflow.log_input)を用いることで、ラン(トレーニング)で用いたデータセットをモデルとともに記録することができます。以下の例では、上と同じirisデータセットをUnity Catalog上のDeltaテーブル`takaakiyayoi_catalog.mlflow.iris`として保存し、これをトレーニングデータとして使用しています。

```py
import mlflow
import pandas as pd
import pyspark.pandas as ps
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestRegressor

# Unity Catalogへのテーブルの書き込み
iris = load_iris()
iris_df = pd.DataFrame(iris.data, columns=iris.feature_names)
# 列名の変更
iris_df.rename(
  columns = {
    'sepal length (cm)':'sepal_length',
    'sepal width (cm)':'sepal_width',
    'petal length (cm)':'petal_length',
    'petal width (cm)':'petal_width'},
  inplace = True
)
iris_df['species'] = iris.target
ps.from_pandas(iris_df).to_table("takaakiyayoi_catalog.mlflow.iris", mode="overwrite")

# Unity Catalogテーブルのロード
dataset = mlflow.data.load_delta(table_name="takaakiyayoi_catalog.mlflow.iris", version="0")
pd_df = dataset.df.toPandas()
X = pd_df.drop("species", axis=1)
y = pd_df["species"]

with mlflow.start_run():
    clf = RandomForestRegressor(n_estimators=100)
    # モデルのトレーニング
    clf.fit(X, y)
    # 入力テーブルの記録
    mlflow.log_input(dataset, "training")
```

結果として作成されるMLflowランを確認すると、**使用されたデータセット**に`takaakiyayoi_catalog.mlflow.iris@v0`と表示されています。テーブル名とバージョン番号`v0`です。Deltaはデータのバージョン管理ができるので、`mlflow.data.load_delta`でデータロードで指定したバージョン番号がそのまま記録されることになります。
![Screenshot 2024-07-17 at 8.54.47.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0574e1ec-759e-8ac2-4781-27a02e2ac7f6.png)

テーブルも確認できます。
![Screenshot 2024-07-17 at 8.55.15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0a4d3ae8-c535-a98a-603b-422486640541.png)

このようにモデルとトレーニングをきちんと管理しながら実験を繰り返すことができます。そして、ある程度のKPI(制度やレーテンシー)をクリアしたら、本格運用に投入という流れになります。そこで、Unity Catalogのモデル管理機能を使うことになります。

# Unity Catalogにおけるモデルとデータのトラッキング

上の例では、モデル自体を記録していなかったので、以下の例では`mlflow.sklearn.log_model`を用いてモデルを記録すると同時に、`registered_model_name="takaakiyayoi_catalog.mlflow.iris_model"`を指定してUnity Catalogに登録するようにします。また、この際に入力サンプルを指定しています。これらはモデルのデプロイに必要になる情報です。詳細に関しては[こちらのマニュアル](https://mlflow.org/docs/latest/model/signatures.html)をご覧ください。

```py
# Unity Catalogにモデルを登録するように設定
mlflow.set_registry_uri("databricks-uc")

# Unity Catalogテーブルのロード
dataset = mlflow.data.load_delta(table_name="takaakiyayoi_catalog.mlflow.iris", version="0")
pd_df = dataset.df.toPandas()
X = pd_df.drop("species", axis=1)
y = pd_df["species"]

# モデル入力のサンプルとしてトレーニングデータセットの最初の行を取得
input_example = X.iloc[[0]]

with mlflow.start_run():
    clf = RandomForestRegressor(n_estimators=100)
    # モデルのトレーニング
    clf.fit(X, y)
    # 入力テーブルの記録
    mlflow.log_input(dataset, "training")

    # モデルを記録し、UCの新バージョンとしてモデルを登録
    mlflow.sklearn.log_model(
        sk_model=clf,
        artifact_path="model",
        # 入力サンプルとその予測アウトプットから自動でシグネチャを推定
        input_example=input_example,
        registered_model_name="takaakiyayoi_catalog.mlflow.iris_model",
    )
```

すると、**使用されたデータセット**に加えて、**登録済みモデル**にモデル名とモデルバージョン`v1`が表示されます。

![Screenshot 2024-07-17 at 9.07.58.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/6591a17c-41c1-8665-c700-f63f2d7d251e.png)

Unity Catalog配下のモデルも確認できます。
![Screenshot 2024-07-17 at 9.09.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0e465945-7a6d-1f69-8b7e-a8ba18999cb6.png)

依存関係(リネージ)も確認できます。
![Screenshot 2024-07-17 at 9.09.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/884f5fc0-7a1f-7267-7754-65de53187aa3.png)
![Screenshot 2024-07-17 at 9.10.32.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7724e1b8-9f06-5dc2-4677-e20db412a726.png)

ご活用ください！

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)