Skip to content

Experiment Tracking with MLflow

Track training runs, compare results, and register models using the PAIS MLflow instance.

MLflow on PAIS

Property Value
Tracking URI https://mlflow.pais.auckland.ac.nz
Authentication Tuakiri SSO (same as other PAIS services)
Artifact storage VAST network storage (persistent)
Backend PostgreSQL

Prerequisites

pip install mlflow torch torchvision

Set the tracking URI: 

export MLFLOW_TRACKING_URI="https://mlflow.pais.auckland.ac.nz"

Basic experiment tracking

import mlflow
import mlflow.pytorch
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])
mlflow.set_experiment("bert_finetune_v1")

with mlflow.start_run(run_name="lr_0.001_batch_32"):
    # Log hyperparameters
    mlflow.log_params({
        "learning_rate": 1e-3,
        "batch_size": 32,
        "epochs": 10,
        "optimizer": "AdamW",
        "model": "bert-base-uncased",
    })

    # Training loop
    for epoch in range(10):
        train_loss = 0.1 * (10 - epoch)  # placeholder
        val_loss = 0.12 * (10 - epoch)
        val_acc = 0.6 + epoch * 0.03

        mlflow.log_metrics({
            "train_loss": train_loss,
            "val_loss": val_loss,
            "val_accuracy": val_acc,
        }, step=epoch)

    # Log the model
    # mlflow.pytorch.log_model(model, "model")

    print(f"Run ID: {mlflow.active_run().info.run_id}")

Comparing runs

import mlflow
import pandas as pd

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])

runs = mlflow.search_runs(
    experiment_names=["bert_finetune_v1"],
    filter_string="metrics.val_accuracy > 0.8",
    order_by=["metrics.val_accuracy DESC"],
    max_results=10,
)

print(runs[["run_id", "params.learning_rate", "params.batch_size",
            "metrics.val_accuracy", "metrics.val_loss"]].to_string())

Model registry

# Register the best run's model
best_run = runs.iloc[0]
model_uri = f"runs:/{best_run.run_id}/model"

mlflow.register_model(model_uri, "bert-finetune-classifier")

# Load a registered model
model = mlflow.pytorch.load_model("models:/bert-finetune-classifier/Production")

See also