Experiment Tracking

---

Overview

Experiment tracking enables reproducibility, comparison, and collaboration in ML development. It captures hyperparameters, metrics, artifacts, and model versions to ensure every experiment can be reproduced and compared.

When to Use This Reference

• Setting up MLflow for experiment tracking
• Implementing Weights & Biases integration
• Creating model registries and versioning
• Comparing experiments and selecting models
• Building custom tracking solutions

When NOT to Use

• Quick one-off experiments without reproducibility needs
• Simple scripts without hyperparameters
• Non-ML projects

---

MLflow Integration

Basic Experiment Tracking

from mlflow.tracking import MlflowClient
from pathlib import Path

class MLflowTracker:
    """MLflow experiment tracking wrapper."""

    def __init__(
        self,
        experiment_name: str,
        tracking_uri: str = "http://localhost:5000",
        artifact_location: str = None,
    ):
        mlflow.set_tracking_uri(tracking_uri)

        # Create or get experiment
        experiment = mlflow.get_experiment_by_name(experiment_name)
        if experiment is None:
            self.experiment_id = mlflow.create_experiment(
                experiment_name,
                artifact_location=artifact_location,
            )
        else:
            self.experiment_id = experiment.experiment_id

        mlflow.set_experiment(experiment_name)
        self.client = MlflowClient()
        self.run = None

    def start_run(
        self,
        run_name: str = None,
        tags: dict = None,
        nested: bool = False,
    ) -> str:
        """Start a new MLflow run."""
        self.run = mlflow.start_run(
            run_name=run_name,
            experiment_id=self.experiment_id,
            nested=nested,
        )

        if tags:
            mlflow.set_tags(tags)

        return self.run.info.run_id

    def end_run(self, status: str = "FINISHED") -> None:
        """End the current run."""
        mlflow.end_run(status=status)
        self.run = None

    def log_params(self, params: dict) -> None:
        """Log hyperparameters."""
        mlflow.log_params(params)

    def log_metrics(self, metrics: dict, step: int = None) -> None:
        """Log metrics with optional step."""
        for key, value in metrics.items():
            mlflow.log_metric(key, value, step=step)

    def log_artifact(self, local_path: str, artifact_path: str = None) -> None:
        """Log file or directory as artifact."""
        mlflow.log_artifact(local_path, artifact_path)

    def log_model(
        self,
        model,
        artifact_path: str,
        registered_model_name: str = None,
        signature=None,
        input_example=None,
    ) -> str:
        """Log model with optional registration."""
        from mlflow.models import infer_signature

        if signature is None and input_example is not None:
            signature = infer_signature(input_example, model.predict(input_example))

        model_info = mlflow.sklearn.log_model(
            model,
            artifact_path=artifact_path,
            registered_model_name=registered_model_name,
            signature=signature,
            input_example=input_example,
        )

        return model_info.model_uri

# Usage example
def train_with_mlflow(
    model,
    X_train,
    y_train,
    X_val,
    y_val,
    params: dict,
):
    """Complete training run with MLflow tracking."""
    tracker = MLflowTracker("my_experiment")

    tracker.start_run(
        run_name=f"run_{params['model_type']}",
        tags={
            "model_type": params["model_type"],
            "dataset_version": "v1.0",
            "author": "ml-team",
        },
    )

    try:
        # Log parameters
        tracker.log_params(params)

        # Train model
        model.fit(X_train, y_train)

        # Evaluate and log metrics
        train_score = model.score(X_train, y_train)
        val_score = model.score(X_val, y_val)

        tracker.log_metrics({
            "train_accuracy": train_score,
            "val_accuracy": val_score,
        })

        # Log model
        model_uri = tracker.log_model(
            model,
            artifact_path="model",
            registered_model_name="my_model",
            input_example=X_train[:5],
        )

        tracker.end_run()
        return model_uri

    except Exception as e:
        tracker.end_run(status="FAILED")
        raise

PyTorch Model Logging

def log_pytorch_model(
    model: torch.nn.Module,
    artifact_path: str,
    registered_model_name: str = None,
    sample_input: torch.Tensor = None,
) -> str:
    """Log PyTorch model with signature inference."""
    from mlflow.models import infer_signature

    # Create signature from sample input
    signature = None
    if sample_input is not None:
        model.eval()
        with torch.no_grad():
            sample_output = model(sample_input)

        signature = infer_signature(
            sample_input.numpy(),
            sample_output.numpy(),
        )

    model_info = mlflow.pytorch.log_model(
        model,
        artifact_path=artifact_path,
        registered_model_name=registered_model_name,
        signature=signature,
    )

    return model_info.model_uri

def load_pytorch_model(model_uri: str, device: str = "cpu") -> torch.nn.Module:
    """Load PyTorch model from MLflow."""
    model = mlflow.pytorch.load_model(model_uri, map_location=device)
    return model

Model Registry Operations

from mlflow.tracking import MlflowClient
from mlflow.entities.model_registry import ModelVersion

class ModelRegistry:
    """MLflow Model Registry wrapper."""

    def __init__(self, tracking_uri: str = "http://localhost:5000"):
        mlflow.set_tracking_uri(tracking_uri)
        self.client = MlflowClient()

    def register_model(
        self,
        model_uri: str,
        name: str,
        tags: dict = None,
        description: str = None,
    ) -> ModelVersion:
        """Register a new model version."""
        result = mlflow.register_model(model_uri, name)

        if tags:
            for key, value in tags.items():
                self.client.set_model_version_tag(name, result.version, key, value)

        if description:
            self.client.update_model_version(
                name,
                result.version,
                description=description,
            )

        return result

    def transition_model_stage(
        self,
        name: str,
        version: str,
        stage: str,
        archive_existing: bool = True,
    ) -> ModelVersion:
        """Transition model to new stage (Staging, Production, Archived)."""
        return self.client.transition_model_version_stage(
            name=name,
            version=version,
            stage=stage,
            archive_existing_versions=archive_existing,
        )

    def get_latest_version(
        self,
        name: str,
        stages: list[str] = None,
    ) -> list[ModelVersion]:
        """Get latest model versions by stage."""
        return self.client.get_latest_versions(name, stages=stages)

    def load_production_model(self, name: str) -> any:
        """Load the production model."""
        model_uri = f"models:/{name}/Production"
        return mlflow.pyfunc.load_model(model_uri)

    def compare_versions(
        self,
        name: str,
        version_a: str,
        version_b: str,
    ) -> dict:
        """Compare two model versions."""
        v_a = self.client.get_model_version(name, version_a)
        v_b = self.client.get_model_version(name, version_b)

        run_a = self.client.get_run(v_a.run_id)
        run_b = self.client.get_run(v_b.run_id)

        return {
            "version_a": {
                "version": version_a,
                "metrics": run_a.data.metrics,
                "params": run_a.data.params,
            },
            "version_b": {
                "version": version_b,
                "metrics": run_b.data.metrics,
                "params": run_b.data.params,
            },
        }

---

Weights & Biases Integration

Basic W&B Tracking

from pathlib import Path

class WandbTracker:
    """Weights & Biases experiment tracking wrapper."""

    def __init__(
        self,
        project: str,
        entity: str = None,
        config: dict = None,
    ):
        self.project = project
        self.entity = entity
        self.config = config
        self.run = None

    def start_run(
        self,
        name: str = None,
        tags: list[str] = None,
        group: str = None,
        job_type: str = "train",
        resume: str = None,
    ) -> wandb.Run:
        """Initialize W&B run."""
        self.run = wandb.init(
            project=self.project,
            entity=self.entity,
            name=name,
            config=self.config,
            tags=tags,
            group=group,
            job_type=job_type,
            resume=resume,
        )
        return self.run

    def log(self, data: dict, step: int = None, commit: bool = True) -> None:
        """Log metrics and data."""
        wandb.log(data, step=step, commit=commit)

    def log_artifact(
        self,
        name: str,
        artifact_type: str,
        path: str,
        metadata: dict = None,
    ) -> wandb.Artifact:
        """Log artifact (model, dataset, etc.)."""
        artifact = wandb.Artifact(
            name=name,
            type=artifact_type,
            metadata=metadata,
        )

        if Path(path).is_dir():
            artifact.add_dir(path)
        else:
            artifact.add_file(path)

        self.run.log_artifact(artifact)
        return artifact

    def log_model(
        self,
        model_path: str,
        name: str,
        metadata: dict = None,
        aliases: list[str] = None,
    ) -> wandb.Artifact:
        """Log model artifact with aliases."""
        artifact = wandb.Artifact(
            name=name,
            type="model",
            metadata=metadata,
        )

        if Path(model_path).is_dir():
            artifact.add_dir(model_path)
        else:
            artifact.add_file(model_path)

        self.run.log_artifact(artifact, aliases=aliases or ["latest"])
        return artifact

    def watch_model(
        self,
        model,
        log: str = "all",
        log_freq: int = 100,
    ) -> None:
        """Watch model for gradient and parameter logging."""
        wandb.watch(model, log=log, log_freq=log_freq)

    def finish(self, exit_code: int = 0) -> None:
        """Finish the run."""
        wandb.finish(exit_code=exit_code)

# Usage with PyTorch
def train_with_wandb(
    model: torch.nn.Module,
    train_loader,
    val_loader,
    config: dict,
):
    """Training with W&B tracking."""
    tracker = WandbTracker(
        project="my-project",
        config=config,
    )

    tracker.start_run(
        name=f"experiment_{config['model_type']}",
        tags=["baseline", config["model_type"]],
        group="hyperparameter_search",
    )

    # Watch model gradients
    tracker.watch_model(model)

    for epoch in range(config["epochs"]):
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            # Training step
            loss = train_step(model, data, target)

            tracker.log({
                "train/loss": loss,
                "train/epoch": epoch,
            })

        # Validation
        val_metrics = evaluate(model, val_loader)
        tracker.log({
            "val/loss": val_metrics["loss"],
            "val/accuracy": val_metrics["accuracy"],
            "epoch": epoch,
        })

    # Save and log model
    torch.save(model.state_dict(), "model.pt")
    tracker.log_model(
        "model.pt",
        name="trained_model",
        metadata={"accuracy": val_metrics["accuracy"]},
        aliases=["latest", "best"],
    )

    tracker.finish()

W&B Sweeps for Hyperparameter Tuning

sweep_config = {
    "method": "bayes",  # bayes, grid, random
    "metric": {
        "name": "val/loss",
        "goal": "minimize",
    },
    "parameters": {
        "learning_rate": {
            "distribution": "log_uniform_values",
            "min": 1e-5,
            "max": 1e-2,
        },
        "batch_size": {
            "values": [16, 32, 64, 128],
        },
        "hidden_size": {
            "values": [128, 256, 512],
        },
        "dropout": {
            "distribution": "uniform",
            "min": 0.1,
            "max": 0.5,
        },
    },
    "early_terminate": {
        "type": "hyperband",
        "min_iter": 3,
    },
}

def sweep_train():
    """Training function for sweep."""
    with wandb.init() as run:
        config = wandb.config

        model = build_model(
            hidden_size=config.hidden_size,
            dropout=config.dropout,
        )

        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=config.learning_rate,
        )

        train_loader = DataLoader(train_dataset, batch_size=config.batch_size)

        for epoch in range(10):
            loss = train_epoch(model, train_loader, optimizer)
            val_loss = evaluate(model, val_loader)

            wandb.log({
                "train/loss": loss,
                "val/loss": val_loss,
                "epoch": epoch,
            })

# Run sweep
sweep_id = wandb.sweep(sweep_config, project="my-project")
wandb.agent(sweep_id, function=sweep_train, count=50)

---

Custom Experiment Tracking

Lightweight Tracker

from datetime import datetime
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import Optional

@dataclass
class Experiment:
    """Experiment metadata and results."""
    experiment_id: str
    name: str
    params: dict
    metrics: dict = field(default_factory=dict)
    artifacts: list = field(default_factory=list)
    tags: dict = field(default_factory=dict)
    start_time: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    end_time: Optional[str] = None
    status: str = "running"

    def to_dict(self) -> dict:
        return asdict(self)

class SimpleTracker:
    """Lightweight file-based experiment tracker."""

    def __init__(self, experiments_dir: str = "./experiments"):
        self.experiments_dir = Path(experiments_dir)
        self.experiments_dir.mkdir(parents=True, exist_ok=True)
        self.current_experiment: Optional[Experiment] = None

    def start_experiment(
        self,
        name: str,
        params: dict,
        tags: dict = None,
    ) -> Experiment:
        """Start a new experiment."""
        experiment_id = str(uuid.uuid4())[:8]

        self.current_experiment = Experiment(