MLflow 3.0

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

有効化

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

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

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

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

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

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

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を使用して、チェックポイントを取得し、それらを精度でランク付けすることができます。

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でも確認できます。

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'}>

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

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、パラメータ、およびメトリクスは、カタログエクスプローラーのモデルバージョンページで確認できます。

# カタログに対する `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を使用して同じ情報を取得することもできます。

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

# モデルバージョンを取得
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

Databricks無料トライアル

Databricks無料トライアル