DatabricksにおけるMLflow 3.0のサポート - ディープラーニング編

こちらの続きです。

https://qiita.com/taka_yayoi/items/8e66db25d99b4de5c85b

こちらのマニュアルに沿ってサンプルノートブックを動かします。

https://docs.databricks.com/aws/ja/mlflow/mlflow-3-install

:::note info
**注意**
執筆時点では[ベータ版](https://docs.databricks.com/aws/ja/release-notes/release-types)です。
:::

# MLflow 3.0

こちらにもまとめていますが、エクスペリメントページにおける**モデル**タブの追加、**Logged Model**エンティティの追加、モデルレジストリUIの改善などが含まれています。特に大きな変更点は、情報アーキテクチャが見直されたことだと思います。これまでは、MLflow**ラン**(トレーニングの実行単位)にモデルやアーティファクト、メトリクス、パラメータなどが紐づけられていました。生成AIのように一つのモデルに繰り返し評価が行われる場合や、ディープラーニングのように複数のモデルチェックポイントができる場合には、この情報アーキテクチャでは困難さが生じていたため、今回の変更で**モデル**に複数のメトリクス、アーティファクトが紐づけられるようになっています。

https://qiita.com/taka_yayoi/items/8ebc58c04f7889438eb1

# 有効化

ワークスペースの**プレビュー**メニューにアクセスし、**MLflow 3**を有効化します。

![Screenshot 2025-04-09 at 21.15.18.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/54771c50-54cf-40c5-9ada-215335bb440a.png)

# ディープラーニングのワークフロー

こちらのサンプルを動かします。

https://docs.databricks.com/aws/ja/mlflow/mlflow3-dl-workflow

このノートブックは最初にモデルのトレーニングジョブを実行し、MLflowランとして追跡します。10エポックごとにモデルのチェックポイントを保存します。各チェックポイントはMLflowの`LoggedModel`として追跡されます。その後、最適なチェックポイントを選択して本番アプリケーションにデプロイできます。

```py
%pip install mlflow  --upgrade --pre torch scikit-learn
dbutils.library.restartPython()
```

```py
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
```

```py
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

import mlflow
import mlflow.pytorch
from mlflow.entities import Dataset

# データを準備するためのヘルパー関数
def prepare_data(df):
    X = torch.tensor(df.iloc[:, :-1].values, dtype=torch.float32)
    y = torch.tensor(df.iloc[:, -1].values, dtype=torch.long)
    return X, y

# 精度を計算するためのヘルパー関数
def compute_accuracy(model, X, y):
    with torch.no_grad():
        outputs = model(X)
        _, predicted = torch.max(outputs, 1)
        accuracy = (predicted == y).sum().item() / y.size(0)
    return accuracy

# 基本的なPyTorch分類器を定義
class IrisClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(IrisClassifier, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# Irisデータセットをロードし、DataFrameを準備
iris = load_iris()
iris_df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
iris_df['target'] = iris.target

# トレーニングデータとテストデータに分割
train_df, test_df = train_test_split(iris_df, test_size=0.2, random_state=42)

# トレーニングデータを準備
train_dataset = mlflow.data.from_pandas(train_df, name="train")
X_train, y_train = prepare_data(train_dataset.df)

# PyTorchモデルを定義し、デバイスに移動
input_size = X_train.shape[1]
hidden_size = 16
output_size = len(iris.target_names)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
scripted_model = IrisClassifier(input_size, hidden_size, output_size).to(device)
scripted_model = torch.jit.script(scripted_model)

# トレーニングジョブを表すランを開始
with mlflow.start_run():
    # トレーニングデータセットをMLflowでロード。トレーニングメトリクスをこのデータセットにリンク
    train_dataset: Dataset = mlflow.data.from_pandas(train_df, name="train")
    X_train, y_train = prepare_data(train_dataset.df)

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(scripted_model.parameters(), lr=0.01)

    for epoch in range(101):
        X_train, y_train = X_train.to(device), y_train.to(device)
        out = scripted_model(X_train)
        loss = criterion(out, y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 10エポックごとにメトリクスを含むチェックポイントをログ
        if epoch % 10 == 0:
            # 新しく作成された各LoggedModelチェックポイントは、その名前、パラメータ、およびステップにリンク
            model_info = mlflow.pytorch.log_model(
                pytorch_model=scripted_model,
                name=f"torch-iris-{epoch}",
                params={
                    "n_layers": 3,
                    "activation": "ReLU",
                    "criterion": "CrossEntropyLoss",
                    "optimizer": "Adam"
                },
                step=epoch,
                input_example=X_train.numpy(),
            )
            # トレーニングデータセットのステップでメトリクスをログし、LoggedModelにリンク
            mlflow.log_metric(
                key="accuracy",
                value=compute_accuracy(scripted_model, X_train, y_train),
                step=epoch,
                model_id=model_info.model_id,
                dataset=train_dataset
            )
```

この例では、1つのMLflow Run（`training_run`）と11のMLflow Logged Models（各チェックポイントごとに1つ、ステップ0、10、…、100）が生成されました。MLflowのUIまたは検索APIを使用して、チェックポイントを取得し、それらを精度でランク付けすることができます。

```py
ranked_checkpoints = mlflow.search_logged_models(output_format="list")
ranked_checkpoints.sort(
    key=lambda model: next((metric.value for metric in model.metrics if metric.key == "accuracy"), float('-inf')),
    reverse=True
)

best_checkpoint: mlflow.entities.LoggedModel = ranked_checkpoints[0]
print(best_checkpoint.metrics[0])
```

```
<Metric: dataset_digest='1f1c13b5', dataset_name='train', key='accuracy', model_id='m-28a5d0d3d0754062955d2b22e0ca55a7', run_id='5912e83e7a394f8aa74e5afb22597ee2', step=80, timestamp=1744235154077, value=0.9833333333333333>
```

UIでも確認できます。

![Screenshot 2025-04-10 at 6.46.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/854e42bb-806d-4275-8c5f-3e317c1aa8f3.png)
![Screenshot 2025-04-10 at 6.47.00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/6c9f7b14-5775-4cff-bea2-c24c4c6966c6.png)
![Screenshot 2025-04-10 at 6.47.09.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/f31cf53f-2789-4979-889e-fb96b326c3f8.png)

```py
best_checkpoint
```
```
<LoggedModel: artifact_location='dbfs:/databricks/mlflow-tracking/3126308776674616/logged_models/m-28a5d0d3d0754062955d2b22e0ca55a7/artifacts', creation_timestamp=1744235146991, experiment_id='3126308776674616', last_updated_timestamp=1744235153950, metrics=[<Metric: dataset_digest='1f1c13b5', dataset_name='train', key='accuracy', model_id='m-28a5d0d3d0754062955d2b22e0ca55a7', run_id='5912e83e7a394f8aa74e5afb22597ee2', step=80, timestamp=1744235154077, value=0.9833333333333333>], model_id='m-28a5d0d3d0754062955d2b22e0ca55a7', model_type='', model_uri='models:/m-28a5d0d3d0754062955d2b22e0ca55a7', name='torch-iris-80', params={'activation': 'ReLU',
 'criterion': 'CrossEntropyLoss',
 'n_layers': '3',
 'optimizer': 'Adam'}, source_run_id='5912e83e7a394f8aa74e5afb22597ee2', status=<LoggedModelStatus.READY: 'READY'>, status_message='', tags={'mlflow.databricks.cluster.id': '0409-213905-yo3le13x-v2n',
 'mlflow.databricks.notebook.commandID': '1744234748534_5981102153045284861_bff301aca7a249ddb13eb14243bddb6b',
 'mlflow.databricks.notebookID': '3126308776674616',
 'mlflow.databricks.notebookPath': '/Users/takaaki.yayoi@databricks.com/20250409_mlflow3.0/mlflow3-dl-example',
 'mlflow.databricks.webappURL': 'https://xxxxxx.cloud.databricks.com',
 'mlflow.databricks.workspaceID': '5099015744649857',
 'mlflow.databricks.workspaceURL': 'https://xxxxxx.cloud.databricks.com',
 'mlflow.source.name': '/Users/takaaki.yayoi@databricks.com/20250409_mlflow3.0/mlflow3-dl-example',
 'mlflow.source.type': 'NOTEBOOK',
 'mlflow.user': 'spark-c27a42b9-687e-4be4-88c7-0f'}>
```

一番精度の悪かったモデルにもアクセスできます。

```py
worst_checkpoint: mlflow.entities.LoggedModel = ranked_checkpoints[-1]
print(worst_checkpoint.metrics)
```
```
[<Metric: dataset_digest='1f1c13b5', dataset_name='train', key='accuracy', model_id='m-e5f945f3fb0440519a29042603d480aa', run_id='5912e83e7a394f8aa74e5afb22597ee2', step=0, timestamp=1744235095488, value=0.3333333333333333>]
```

最適なチェックポイントモデルを選択した後、そのモデルをモデルレジストリに登録します。モデルID、パラメータ、およびメトリクスは、カタログエクスプローラーのモデルバージョンページで確認できます。

```py
# カタログに対する `USE CATALOG` 権限、およびスキーマに対する `USE SCHEMA` 権限が必要です。
# 必要に応じて、ここでカタログとスキーマの名前を変更してください。

CATALOG = "takaakiyayoi_catalog"
SCHEMA = "mlflow"
MODEL = "dl_model"
MODEL_NAME = f"{CATALOG}.{SCHEMA}.{MODEL}"

uc_model_version = mlflow.register_model(f"models:/{best_checkpoint.model_id}", name=MODEL_NAME)
```
```
Successfully registered model 'takaakiyayoi_catalog.mlflow.dl_model'.
Created version '1' of model 'takaakiyayoi_catalog.mlflow.dl_model'.
```

モデルバージョンページでモデルバージョンとすべての集中化されたパフォーマンスデータをUnity Catalogで確認できます。次のセルに示すように、APIを使用して同じ情報を取得することもできます。

この場合もモデルのメトリクスやパラメータが確認できるようになっています。

![Screenshot 2025-04-10 at 6.47.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/af36ce06-e092-477e-9237-e45a733009c9.png)

```py
# モデルバージョンを取得
from mlflow import MlflowClient
client = MlflowClient()
model_version = client.get_model_version(name=MODEL_NAME, version=uc_model_version.version)
print(model_version)
```

```
<ModelVersion: aliases=[], creation_timestamp=1744235261854, current_stage=None, deployment_job_state=<ModelVersionDeploymentJobState: current_task_name='', job_id='', job_state='DEPLOYMENT_JOB_CONNECTION_STATE_UNSPECIFIED', run_id='', run_state='DEPLOYMENT_JOB_RUN_STATE_UNSPECIFIED'>, description='', last_updated_timestamp=1744235264372, metrics=[<Metric: dataset_digest='1f1c13b5', dataset_name='train', key='accuracy', model_id='m-28a5d0d3d0754062955d2b22e0ca55a7', run_id='5912e83e7a394f8aa74e5afb22597ee2', step=80, timestamp=1744235154077, value=0.9833333333333333>], model_id='m-28a5d0d3d0754062955d2b22e0ca55a7', name='takaakiyayoi_catalog.mlflow.dl_model', params=[<LoggedModelParameter: key='activation', value='ReLU'>,
 <LoggedModelParameter: key='criterion', value='CrossEntropyLoss'>,
 <LoggedModelParameter: key='n_layers', value='3'>,
 <LoggedModelParameter: key='optimizer', value='Adam'>], run_id='5912e83e7a394f8aa74e5afb22597ee2', run_link=None, source='models:/m-28a5d0d3d0754062955d2b22e0ca55a7', status='READY', status_message='', tags={}, user_id='takaaki.yayoi@databricks.com', version='1'>
```



### はじめてのDatabricks

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

### Databricks無料トライアル

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