こちらの続きです。
こちらのマニュアルに沿ってサンプルノートブックを動かします。
注意
執筆時点ではベータ版です。
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'>