MLflow Manual Logging Tutorial

1. Environment Setup

# Core MLflow imports
import mlflow
import mlflow.pytorch

from datamint import Api
from datamint.mlflow import set_project

api = Api()

2. Project Configuration

Before starting any MLflow operations, we must set the active Datamint project. This ensures all experiments are properly associated with your project.

IMPORTANT: Replace 'BoneSeg' with your actual project name!

# Replace this with your actual project name
PROJECT_NAME = 'FracAtlas'
project = set_project(PROJECT_NAME)

3. Logging Metrics and Model to MLflow

Now we’ll log the pre-computed metrics and the trained model to MLflow. This demonstrates how to log results after training is complete.

Key MLflow Functions Used:

• mlflow.start_run(): Start a new MLflow run, representing the past training session

• mlflow.log_params(): Log hyperparameters

• mlflow.log_metric(): Log a metric value at a specific step (epoch)

• mlflow.pytorch.log_model(): Save the trained model

# Assume these are your pre-computed metrics from training
# Replace these with your actual training results
# You can also log per-epoch metrics if you saved them
# For example, if you saved metrics for each epoch during training:
epoch_metrics = [
    {"train/loss": 0.653, "train/dice_score": 0.1, "val/loss": 0.621, "val/dice_score": 0.09},  # Epoch 0
    {"train/loss": 0.512, "train/dice_score": 0.2, "val/loss": 0.534, "val/dice_score": 0.2},  # Epoch 1
    {"train/loss": 0.423, "train/dice_score": 0.25, "val/loss": 0.467, "val/dice_score": 0.22},  # Epoch 2
    {"train/loss": 0.356, "train/dice_score": 0.3, "val/loss": 0.398, "val/dice_score": 0.33},  # Epoch 3
    {"train/loss": 0.298, "train/dice_score": 0.33, "val/loss": 0.345, "val/dice_score": 0.32},  # Epoch 4
    {"train/loss": 0.267, "train/dice_score": 0.36, "val/loss": 0.328, "val/dice_score": 0.35},  # Epoch 5
    {"train/loss": 0.251, "train/dice_score": 0.5, "val/loss": 0.319, "val/dice_score": 0.44},  # Epoch 6
    {"train/loss": 0.242, "train/dice_score": 0.6, "val/loss": 0.315, "val/dice_score": 0.5},  # Epoch 7
    {"train/loss": 0.237, "train/dice_score": 0.62, "val/loss": 0.313, "val/dice_score": 0.55},  # Epoch 8
    {"train/loss": 0.234, "train/dice_score": 0.75, "val/loss": 0.312, "val/dice_score": 0.64},  # Epoch 9
]

best_metrics = {
    "val/best/loss": 0.312,
    "val/best/dice_score": 0.64,
}

mlflow.set_experiment('nnUnet_training1')
# Start MLflow run
with mlflow.start_run(run_name="log_metrics_run") as run:
    print(f"🚀 MLflow Run ID: {run.info.run_id}")
    print(f"Experiment ID: {run.info.experiment_id}")
    print()

    # Log hyperparameters
    mlflow.log_params({
        "learning_rate": 1e-4,
        "batch_size": 16,
        "num_epochs": 10,
        "model_architecture": "nnUnet1",
    })
    print("✅ Hyperparameters logged to MLflow")
    print()

    # Log per-epoch metrics (if available)
    print("Logging per-epoch metrics...")
    for epoch, metrics in enumerate(epoch_metrics):
        mlflow.log_metrics({
            "train/loss": metrics["train/loss"],
            "train/accuracy": metrics["train/dice_score"],
            "val/loss": metrics["val/loss"],
            "val/accuracy": metrics["val/dice_score"],
        }, step=epoch)
    print(f"✅ Logged metrics for {len(epoch_metrics)} epochs")
    print()

    # Log final/best metrics
    print("Logging final metrics...")
    mlflow.log_metrics(best_metrics)
    print("✅ Final metrics logged")
    print()

print("=" * 60)
print("Logging complete!")
print(f"View results at: {api.projects.get_by_name(PROJECT_NAME).url}")
print("=" * 60)