Experiment Tracking — AI Tutorials — Subhendu Datta Bhowmik

The Experiment Tracking Problem

A typical ML project involves dozens or hundreds of experiments: different model architectures, hyperparameter values, feature sets, and preprocessing strategies. Without systematic tracking:

You can't reproduce the model in production
You forget which configuration achieved your best metric
You repeat experiments you've already run
Collaborators can't see what you've tried

Experiment tracking tools solve this by automatically recording every run: code version, parameters, metrics, and output artifacts.

The Three Pillars of Experiment Tracking

Pillar	What It Captures
Parameters	Hyperparameters, config values, flags
Metrics	Loss curves, evaluation scores, custom KPIs
Artifacts	Model files, plots, datasets, feature importance

MLflow

MLflow is the most widely deployed open-source experiment tracking platform. It consists of four components:

MLflow Tracking — log and query experiments
MLflow Projects — package code for reproducible runs
MLflow Models — standard model packaging format
MLflow Model Registry — lifecycle management (staging → production)

It can run locally, on a self-hosted server, or on cloud managed services (Databricks, Azure ML).

MLflow Tracking — Complete Examplepython

import mlflow
import mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import f1_score, roc_auc_score
from sklearn.model_selection import train_test_split
import pandas as pd

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("fraud-detection")

df = pd.read_parquet("data/features/train.parquet")
X, y = df.drop("is_fraud", axis=1), df["is_fraud"]
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, stratify=y)

params = {
    "n_estimators": 200,
    "learning_rate": 0.05,
    "max_depth": 4,
    "subsample": 0.8,
}

with mlflow.start_run(run_name="gbt-baseline"):
    # Log parameters
    mlflow.log_params(params)

    # Train
    model = GradientBoostingClassifier(**params)
    model.fit(X_train, y_train)

    # Log metrics
    preds = model.predict(X_val)
    probs = model.predict_proba(X_val)[:, 1]
    mlflow.log_metrics({
        "f1": f1_score(y_val, preds),
        "auc": roc_auc_score(y_val, probs),
    })

    # Log model (with input schema)
    mlflow.sklearn.log_model(
        model,
        artifact_path="model",
        input_example=X_train.head(5),
    )

MLflow Model Registry

The Model Registry adds lifecycle management on top of tracking. Models move through stages:

None → Staging → Production → Archived

This gives you a controlled promotion process: a data scientist registers a candidate model, an engineer validates it, and an ML engineer promotes it to Production when ready.

Registering and Promoting a Modelpython

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

# Register a model from a completed run
run_id = "abc123def456"
model_uri = f"runs:/{run_id}/model"

model_version = mlflow.register_model(
    model_uri=model_uri,
    name="fraud-detector",
)
print(f"Model version: {model_version.version}")

# Transition to staging
client.transition_model_version_stage(
    name="fraud-detector",
    version=model_version.version,
    stage="Staging",
    archive_existing_versions=False,
)

# After validation, promote to production
client.transition_model_version_stage(
    name="fraud-detector",
    version=model_version.version,
    stage="Production",
    archive_existing_versions=True,  # Archive old production model
)

# Load the current production model
prod_model = mlflow.sklearn.load_model("models:/fraud-detector/Production")

Weights & Biases (W&B)

W&B is the preferred platform for deep learning teams. It offers richer visualizations than MLflow, built-in sweeps (hyperparameter search), and collaborative features.

Key features:

Runs: log metrics, images, audio, video, tables
Sweeps: distributed hyperparameter search (Bayesian, random, grid)
Artifacts: versioned datasets and models with lineage tracking
Reports: shareable analysis documents with embedded charts

W&B Tracking with Sweepspython

import wandb
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import f1_score

# Define hyperparameter search space
sweep_config = {
    "method": "bayes",
    "metric": {"name": "f1", "goal": "maximize"},
    "parameters": {
        "n_estimators": {"values": [100, 200, 300]},
        "learning_rate": {"min": 0.01, "max": 0.2},
        "max_depth": {"values": [3, 4, 5, 6]},
    },
}

def train_sweep():
    with wandb.init() as run:
        cfg = run.config

        model = GradientBoostingClassifier(
            n_estimators=cfg.n_estimators,
            learning_rate=cfg.learning_rate,
            max_depth=cfg.max_depth,
        )
        model.fit(X_train, y_train)

        f1 = f1_score(y_val, model.predict(X_val))
        wandb.log({"f1": f1})

# Create and run the sweep
sweep_id = wandb.sweep(sweep_config, project="fraud-detection")
wandb.agent(sweep_id, function=train_sweep, count=30)

Neptune

Neptune positions itself between MLflow and W&B: self-hosted or cloud, strong Python SDK, and excellent support for logging custom data types (DataFrames, plots, HTML).

Best for: teams that want flexibility in what they log and a clean query API for experiment comparison.

Neptune Run Loggingpython

import neptune
from neptune.utils import stringify_unsupported

run = neptune.init_run(
    project="my-org/fraud-detection",
    api_token="YOUR_API_TOKEN",
)

# Log parameters
run["params"] = {
    "model": "GBT",
    "n_estimators": 200,
    "learning_rate": 0.05,
}

# Log metrics during training (per epoch/iteration)
for epoch, (train_loss, val_f1) in enumerate(training_loop()):
    run["train/loss"].append(train_loss)
    run["val/f1"].append(val_f1)

# Log a DataFrame as a table
run["feature_importance"].upload(
    neptune.types.File.as_html(importance_df.to_html())
)

# Log model artifact
run["model/best"].upload("models/best_model.pkl")

run.stop()

Choosing a Tracker

Criterion	MLflow	W&B	Neptune
Deployment	Self-hosted / Databricks	Cloud or self-hosted	Cloud or self-hosted
License	Apache 2.0 (free)	Free tier + paid	Free tier + paid
Deep learning support	Good	Excellent	Good
Model Registry	Built-in	Via Artifacts	Via Model Registry
Hyperparameter sweeps	Via plugins	Built-in Sweeps	Built-in HPO
Best for	Enterprise, any framework	Deep learning, research	Flexibility, metadata

Rule of thumb: Use MLflow if you're on Databricks or need fully open-source. Use W&B for deep learning research. Use Neptune if you want fine-grained control over what you log.

Knowledge check

What is the purpose of an MLflow Model Registry?

Summary

Experiment tracking captures parameters, metrics, and artifacts for every training run
MLflow is the most widely used open-source option — Tracking + Model Registry covers most workflows
W&B excels at deep learning with rich visualizations and built-in hyperparameter sweeps
Neptune offers flexibility and a clean SDK for logging custom data types
The Model Registry (MLflow) or Artifacts (W&B) manages the model lifecycle from staging to production
Start logging from day one — retroactively adding tracking is painful

Next: feature stores with Feast, Tecton, and Hopsworks.

← MLOps & AI Engineering