ml-experiment-tracker

ML Experiment Tracker

This skill provides guidance for systematic machine learning experimentation with proper tracking, versioning, and reproducibility practices.

Core Competencies

• Experiment Tracking: MLflow, Weights & Biases (wandb), Neptune, Comet
• Data Versioning: DVC, Delta Lake, LakeFS
• Model Registry: Version control for trained models
• Reproducibility: Environment, code, data, and hyperparameter tracking

Experiment Tracking Fundamentals

What to Track

Every experiment should log:

Category	Items	Why
Code	Git commit hash, branch, diff	Reproduce exact code state
Data	Dataset version, hash, lineage	Know which data was used
Environment	Python version, dependencies, hardware	Reproduce runtime
Hyperparameters	All config values	Understand what changed
Metrics	Loss, accuracy, custom metrics	Compare performance
Artifacts	Models, plots, predictions	Preserve outputs

Experiment Organization

project/
├── experiments/
│   ├── baseline/           # Initial experiments
│   ├── feature-engineering/ # Data improvements
│   ├── architecture/       # Model changes
│   └── hyperparameter/     # Tuning runs
├── data/
│   ├── raw/               # Original data (versioned)
│   ├── processed/         # Cleaned data
│   └── features/          # Feature store
└── models/
    ├── staging/           # Candidates
    └── production/        # Deployed models

MLflow Patterns

Basic Experiment Logging

import mlflow

# Set experiment (creates if not exists)
mlflow.set_experiment("my-classification-project")

with mlflow.start_run(run_name="baseline-v1"):
    # Log parameters
    mlflow.log_param("learning_rate", 0.01)
    mlflow.log_param("batch_size", 32)
    mlflow.log_param("epochs", 100)

    # Training loop
    for epoch in range(epochs):
        train_loss = train_epoch(model, train_loader)
        val_loss, val_acc = evaluate(model, val_loader)

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

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

    # Log artifacts (plots, configs)
    mlflow.log_artifact("confusion_matrix.png")
    mlflow.log_artifact("config.yaml")

Model Registry Workflow

┌──────────────┐    ┌──────────────┐    ┌──────────────┐
│   Training   │───▶│   Staging    │───▶│  Production  │
│    Runs      │    │   Review     │    │   Deployed   │
└──────────────┘    └──────────────┘    └──────────────┘
      │                    │                    │
      ▼                    ▼                    ▼
   Candidate           Validated           Monitored
   Models              Models              Models

Stages:

• None: Just logged, not registered
• Staging: Candidate for production
• Production: Active serving
• Archived: Historical reference

Weights & Biases Patterns

Project Structure

import wandb

# Initialize with config
config = {
    "learning_rate": 0.01,
    "architecture": "ResNet50",
    "dataset": "imagenet-subset",
    "epochs": 100
}

run = wandb.init(
    project="image-classification",
    group="architecture-experiments",  # Group related runs
    tags=["baseline", "resnet"],
    config=config,
    notes="Testing ResNet50 baseline on subset"
)

# Training with automatic logging
for epoch in range(config["epochs"]):
    metrics = train_and_eval(model, train_loader, val_loader)
    wandb.log(metrics)

    # Log media
    wandb.log({"predictions": wandb.Image(pred_grid)})
    wandb.log({"confusion_matrix": wandb.plot.confusion_matrix(...)})

wandb.finish()

Hyperparameter Sweeps

# sweep_config.yaml
program: train.py
method: bayes  # or grid, random
metric:
  name: val_accuracy
  goal: maximize
parameters:
  learning_rate:
    distribution: log_uniform_values
    min: 0.0001
    max: 0.1
  batch_size:
    values: [16, 32, 64, 128]
  optimizer:
    values: ["adam", "sgd", "adamw"]
early_terminate:
  type: hyperband
  min_iter: 10

DVC for Data Versioning

Setup and Usage

# Initialize DVC in git repo
dvc init

# Track large files
dvc add data/training.csv
git add data/training.csv.dvc data/.gitignore
git commit -m "Add training data v1"

# Push to remote storage
dvc remote add -d storage s3://bucket/dvc
dvc push

# Create pipeline
dvc run -n preprocess \
    -d src/preprocess.py -d data/raw \
    -o data/processed \
    python src/preprocess.py

# Reproduce pipeline
dvc repro

DVC Pipeline Definition

# dvc.yaml
stages:
  preprocess:
    cmd: python src/preprocess.py
    deps:
      - src/preprocess.py
      - data/raw/
    outs:
      - data/processed/

  train:
    cmd: python src/train.py
    deps:
      - src/train.py
      - data/processed/
    params:
      - train.epochs
      - train.learning_rate
    outs:
      - models/model.pkl
    metrics:
      - metrics.json:
          cache: false

Reproducibility Checklist

Code Reproducibility

• Pin git commit for each experiment
• Track uncommitted changes (git diff)
• Version control notebooks (nbstripout)
• Document manual steps

Environment Reproducibility

• Lock dependencies (pip freeze, poetry.lock)
• Specify Python version
• Document CUDA/GPU requirements
• Use containers for full isolation

Data Reproducibility

• Version datasets with DVC or similar
• Document data collection process
• Track preprocessing steps
• Save train/val/test split indices

Training Reproducibility

• Set random seeds (Python, NumPy, PyTorch/TF)
• Log all hyperparameters
• Save model checkpoints
• Document non-deterministic operations

Best Practices

Naming Conventions

experiment: {project}-{objective}
run: {date}-{description}-{variant}
model: {architecture}-{dataset}-{version}

Examples:
experiment: fraud-detection-baseline
run: 2024-01-15-xgboost-tuning-lr001
model: xgboost-transactions-v2.3.1

Comparison Dashboards

Track these metrics for model comparison:

• Primary metric (what you optimize)
• Secondary metrics (constraints)
• Resource usage (training time, memory)
• Inference performance (latency, throughput)

Experiment Documentation

Each significant experiment should document:

1. Hypothesis: What change and expected outcome
2. Method: What was actually done
3. Results: Metrics and observations
4. Conclusions: What was learned, next steps

References

• references/mlflow-setup.md - MLflow installation and configuration
• references/wandb-patterns.md - Advanced W&B features and sweeps
• references/reproducibility-checklist.md - Detailed reproducibility guide