Learn AI Series (#119) - Experiment Tracking and Reproducibility

avatar

Learn AI Series (#119) - Experiment Tracking and Reproducibility

variant-b-09-blue.png

What will I learn

  • You will learn why an ML experiment you cannot repeat is not really an experiment at all -- it's an anecdote;
  • MLflow for experiment tracking: logging parameters, metrics, and artifacts with three humble function calls;
  • Weights & Biases for richer visualization, team collaboration, and (almost free) hyperparameter sweeps;
  • Hydra for configuration management that scales from a laptop notebook to a cluster grid-search without rewriting a line;
  • Docker for ML: freezing the entire software environment so "works on my machine" stops being a curse;
  • and random seeds, determinism, and the uncomfortable truth that true reproducibility is harder than random.seed(42) makes it look.

Requirements

  • A working modern computer running macOS, Windows or Ubuntu;
  • An installed Python 3(.11+) distribution with PyTorch -- today is mostly plumbing and discipline, so no beefy GPU required;
  • You've been through episode #34 (ML engineering) and episode #43 (PyTorch data and training), and ideally episode #118 (data engineering), because today we bolt the tracking layer onto everything we've already built.

Difficulty

  • Beginner

Curriculum (of the Learn AI Series):

Learn AI Series (#119) - Experiment Tracking and Reproducibility

"I got 94% accuracy last Tuesday. Which hyperparameters was that? What data version? Which random seed? ... No idea."

If that little monologue made you wince in recognition, congratulations -- you are a real ML practitioner, and this episode is aimed squarely at you. Because ML development is inherently experimental. You don't write one model and ship it; you try thirty configurations, squint at the results, pick a winner, and then -- three weeks later -- somebody asks you to reproduce that winner and you discover, to your quiet horror, that you have no idea how you got it. Today we fix that, permanently.

Last episode (#118) we treated data as a verb: the pipes that move it, validate it, and version it. We even name-dropped reproducibility and promised to come back to it. Well, here we are. Data versioning was one leg of the reproducibility stool; today we build the other three -- experiment tracking, configuration management, and environment freezing -- and screw the whole thing together so it actually holds your weight.

Solutions to Episode #118 Exercises

As promised, let's clear last episode's three tasks before we open any new cans of worms. They all build on DataValidator, StreamProcessor and SimpleTabularSynthesizer from episode #118, so I'm assuming those are imported and in scope.

Exercise 1 asked you to upgrade validate_batch so that, on top of the valid/invalid split, it tells you which field is doing the most damage -- the difference between a monitor that whimpers "something's wrong" and one that points and says "go look at price":

from collections import Counter

class DataValidator:
    # __init__ and validate_record are exactly as in episode #118.

    def validate_batch(self, records):
        valid, invalid = [], []
        field_offenders = Counter()

        for record in records:
            ok, errors = self.validate_record(record)
            if ok:
                valid.append(record)
            else:
                invalid.append((record, errors))
                # every error string is "fieldname: reason", so the field
                # name is everything before the first colon.
                for err in errors:
                    field = err.split(":", 1)[0]
                    field_offenders[field] += 1

        rejection_rate = len(invalid) / max(len(records), 1)
        if rejection_rate > 0.05:
            print(f"WARNING: {rejection_rate:.1%} rejection rate -- investigate upstream!")

        total = sum(field_offenders.values())
        if total:
            worst_field, worst_count = field_offenders.most_common(1)[0]
            print(f"Top offender: '{worst_field}' -> {worst_count / total:.0%} of all rejections")

        return valid, invalid

The key insight I wanted you to feel: the rejection rate and the rejection cause are two completely different pieces of information, and only the second one is actionable at 3 AM. "15% of records rejected" tells you to panic. "83% of those rejections were price out of range" tells you a new pricing service is shipping cents where the old one shipped dollars. One number wakes you up; the other lets you go back to sleep. Notice I lean on the fact that every error string starts with fieldname: -- that's a tiny contract between validate_record and validate_batch, and it's worth keeping deliberately, because the moment those two disagree, this tally quietly lies to you.

Exercise 2 was the meaty one: give StreamProcessor a genuine time-windowed feature -- "clicks in the last 60 seconds" -- using event timestamps rather than a running total that grows forever:

from collections import defaultdict, deque

class StreamProcessor:
    """Now emitting a real sliding-window feature, not a lifetime counter."""
    WINDOW_SECONDS = 60

    def __init__(self):
        # one deque of recent click-timestamps per user
        self.user_clicks = defaultdict(deque)

    def process_event(self, event):
        user_id = event["user_id"]
        now = event["timestamp"]
        clicks = self.user_clicks[user_id]

        clicks.append(now)
        # evict everything that has fallen out of the trailing window
        cutoff = now - self.WINDOW_SECONDS
        while clicks and clicks[0] < cutoff:
            clicks.popleft()

        return {
            "user_id": user_id,
            "clicks_last_60s": len(clicks),
        }

A deque is exactly the right tool here: append on the right, pop the stale ones off the left, and the length is your feature. No re-scanning a list, no unbounded memory growth -- old clicks physically leave the structure once they age out. Now the paragraph I actually cared about, on out-of-order events: in the real world your events do NOT arrive in tidy timestamp order. A phone loses signal in a tunnel, buffers three clicks, and dumps them thirty seconds late; a slow Kafka partition delivers events behind a fast one. My tidy while clicks[0] < cutoff eviction assumes timestamps only ever go up, so a late event with an old timestamp gets appended to the right of newer ones and my cutoff logic breaks -- I might evict a valid click, or count a stale one. The industrial defence is a watermark: you declare "I will accept events up to, say, 10 seconds late, and anything later than that gets dropped or shunted to a side-channel." You trade a little completeness for bounded, sane state. Streaming systems like Flink and Beam are built almost entirely around this one ugly truth: time is a lie, and you must decide how much of a lie you'll tolerate.

Exercise 3 asked you to actually measure the synthetic tax -- train one classifier on real data, one on synthetic data, and score BOTH on held-out real data:

import numpy as np
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# SimpleTabularSynthesizer comes straight from episode #118.

data = load_breast_cancer()
X, y = data.data, data.target
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42, stratify=y
)

# 1. The honest baseline: train on REAL, score on REAL.
real_clf = LogisticRegression(max_iter=5000).fit(X_train, y_train)
real_acc = accuracy_score(y_test, real_clf.predict(X_test))

# 2. Synthesize a same-sized training set, one class at a time so labels survive.
X_parts, y_parts = [], []
for label in np.unique(y_train):
    subset = X_train[y_train == label]
    synth = SimpleTabularSynthesizer()
    synth.fit({i: subset[:, i] for i in range(subset.shape[1])})
    gen = synth.generate(len(subset))
    X_parts.append(np.column_stack([gen[i] for i in range(subset.shape[1])]))
    y_parts.append(np.full(len(subset), label))

X_synth = np.vstack(X_parts)
y_synth = np.concatenate(y_parts)

# 3. Train on SYNTHETIC, but score on the SAME held-out REAL test set.
synth_clf = LogisticRegression(max_iter=5000).fit(X_synth, y_synth)
synth_acc = accuracy_score(y_test, synth_clf.predict(X_test))

print(f"real-trained accuracy:      {real_acc:.3f}")
print(f"synthetic-trained accuracy: {synth_acc:.3f}")
print(f"the synthetic tax:          {real_acc - synth_acc:.3f}")

The gap between those two numbers is the "synthetic tax" -- the price you pay for pretending. And it's usually painfully visible with our toy synthesizer, precisely because (as I warned in #118) it models each column independently and cheerfully shreds the correlations between them. On the breast-cancer set, "worst radius" and "worst area" are almost the same fact told twice; the real classifier leans on that relationship, and the synthetic one never learned it existed. The whole point of measuring the tax on real held-out data is that it's the only honest judge -- score synthetic-trained models on synthetic test data and you'll get a gorgeous, meaningless number. Swap in CTGAN and the gap shrinks, because it models the joint distribution. It never hits zero, though. It never should.

Right -- solutions filed. On to tracking.

Why tracking matters (git is not enough)

In traditional software, life is relatively kind: you write code, it works or it doesn't, and git faithfully records every change. You can git bisect your way back to the exact commit that broke things. ML laughs at this comfort, because a model's behaviour depends on a whole galaxy of things git never sees: hyperparameters, the version of the training data, random seeds, the GPU it ran on, the exact package versions, and dozens of metric curves wiggling over time.

Think about the numbers for a second. A single training run might chew through 20 hyperparameters, emit 50 metrics across its epochs, and spit out a handful of artifacts (the model weights, a confusion matrix, maybe a few sample predictions). Now multiply that by the 100+ runs a real project accumulates, and you are drowning in un-labelled, un-searchable state. Having said that, the fix is not heroic discipline or a giant spreadsheet you swear you'll keep updated (you won't). The fix is tooling that captures all of it automatically, every run, without you having to remember. That's what experiment tracking is.

MLflow: the open-source standard

MLflow is the most widely used experiment tracking tool, and for good reason -- it's free, open-source, self-hosted, and refreshingly un-clever. It logs everything about a run (parameters, metrics, artifacts, code version) and gives you a web UI to compare runs side by side. The mental model is three verbs: log your params at the start, log your metrics during, log your model at the end.

import mlflow
import torch
import torch.nn as nn

def train_with_mlflow(config, train_loader, val_loader):
    """A PyTorch training run with full experiment tracking bolted on."""
    mlflow.set_experiment("image_classifier")

    with mlflow.start_run(run_name=f"lr{config['lr']}_bs{config['batch_size']}"):
        # 1. Log every hyperparameter at the start.
        mlflow.log_params(config)
        mlflow.log_param("pytorch_version", torch.__version__)

        model = nn.Sequential(
            nn.Linear(784, config["hidden_size"]),
            nn.ReLU(),
            nn.Dropout(config["dropout"]),
            nn.Linear(config["hidden_size"], 10),
        )
        optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
        criterion = nn.CrossEntropyLoss()

        for epoch in range(config["epochs"]):
            model.train()
            total_loss, correct, total = 0.0, 0, 0
            for batch_x, batch_y in train_loader:
                optimizer.zero_grad()
                output = model(batch_x)
                loss = criterion(output, batch_y)
                loss.backward()
                optimizer.step()

                total_loss += loss.item()
                correct += (output.argmax(1) == batch_y).sum().item()
                total += len(batch_y)

            # 2. Log metrics DURING training, tagged with the epoch (step).
            mlflow.log_metrics({
                "train_loss": total_loss / len(train_loader),
                "train_accuracy": correct / total,
            }, step=epoch)

            val_acc = evaluate(model, val_loader)
            mlflow.log_metric("val_accuracy", val_acc, step=epoch)
            print(f"epoch {epoch}: val_acc={val_acc:.4f}")

        # 3. Log the trained model itself as an artifact.
        mlflow.pytorch.log_model(model, "model")
        return val_acc

By default MLflow scribbles all of this into a local mlruns/ directory, which is perfect for solo work -- no server, no account, no fuss. When you're ready for a team, you point it at a remote tracking server and the same three verbs now log to a shared place everyone can browse. Fire up mlflow ui in a terminal and you get a sortable, filterable table of every run you've ever done: sort by val_accuracy descending, and there's your Tuesday-94%-model, with every hyperparameter that produced it sitting right there in the row. No more archaeology.

Weights & Biases: tracking with a nicer view

Weights & Biases (everyone just says "wandb") plays in the same space but leans hard into live dashboards and team collaboration. The API will feel familiar after MLflow, but the visualizations are genuinely a step up -- real-time charts, automatic system metrics (GPU temperature, memory), and gradient histograms if you ask nicely.

import wandb
import torch.nn as nn

def train_with_wandb(config, train_loader, val_loader):
    run = wandb.init(
        project="image-classifier",
        config=config,
        name=f"lr{config['lr']}_hs{config['hidden_size']}",
    )

    model = nn.Sequential(
        nn.Linear(784, config["hidden_size"]),
        nn.ReLU(),
        nn.Dropout(config["dropout"]),
        nn.Linear(config["hidden_size"], 10),
    )
    # watch() auto-logs gradients and parameter distributions -- great for
    # spotting exploding/vanishing gradients from episode #40.
    wandb.watch(model, log="all", log_freq=100)

    optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
    criterion = nn.CrossEntropyLoss()

    for epoch in range(config["epochs"]):
        model.train()
        for batch_x, batch_y in train_loader:
            optimizer.zero_grad()
            loss = criterion(model(batch_x), batch_y)
            loss.backward()
            optimizer.step()
            wandb.log({"batch_loss": loss.item()})

        wandb.log({"epoch": epoch, "val_accuracy": evaluate(model, val_loader)})

    wandb.finish()

Where wandb really earns its keep is hyperparameter sweeps. You describe a search space, hand it your training function, and it orchestrates the whole grid (or, better, a Bayesian search that spends its budget where the accuracy actually lives):

sweep_config = {
    "method": "bayes",   # not brute-force grid -- let it be smart
    "metric": {"name": "val_accuracy", "goal": "maximize"},
    "parameters": {
        "lr":          {"min": 1e-5, "max": 1e-2, "distribution": "log_uniform_values"},
        "hidden_size": {"values": [128, 256, 512]},
        "dropout":     {"min": 0.0, "max": 0.5},
        "batch_size":  {"values": [32, 64, 128]},
    },
}

# sweep_id = wandb.sweep(sweep_config, project="image-classifier")
# wandb.agent(sweep_id, train_with_wandb, count=50)

So which do you pick? Honestly, for a solo learner or an open-source project, MLflow is the safe default -- self-hosted, zero cost, no vendor lock-in, your data stays your data. wandb is a hosted service (free for personal use, paid for teams) whose collaboration and sweep features are lovely enough that many companies happily pay. They are not mutually exclusive either -- I've seen setups log to both. The thing that matters is not which tool; it's that you use one, consistently, from run number one. The worst tracking tool is the one you keep meaning to set up "later".

Hydra: configuration that scales

Here's a pattern I want you to learn to hate: hyperparameters scattered as magic numbers across three scripts and a notebook. lr = 3e-4 here, batch_size = 64 there, a dropout buried in a class constructor. It doesn't scale, it isn't reproducible, and it's murder to sweep over. Hydra gives you structured, composable configuration in plain YAML files instead:

# config/train.yaml
model:
  hidden_size: 256
  dropout: 0.2
  n_layers: 3

training:
  lr: 3e-4
  batch_size: 64
  epochs: 50
  weight_decay: 1e-5

data:
  train_path: data/train.parquet
  val_path: data/val.parquet
  max_samples: null   # null = use all of it

And in code, Hydra hands you that whole tree as a single object:

import hydra
from omegaconf import DictConfig

@hydra.main(config_path="config", config_name="train", version_base=None)
def train(cfg: DictConfig):
    print(f"lr={cfg.training.lr}, hidden={cfg.model.hidden_size}")

    model = build_model(
        hidden_size=cfg.model.hidden_size,
        dropout=cfg.model.dropout,
        n_layers=cfg.model.n_layers,
    )
    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=cfg.training.lr,
        weight_decay=cfg.training.weight_decay,
    )
    # ... the rest of your training loop ...

if __name__ == "__main__":
    train()

The magic is on the command line. You override any value without touching the file, compose whole config groups, and -- my favourite -- kick off a grid search with one flag:

python train.py training.lr=1e-4                      # override a single value
python train.py model=large                           # swap in an entire config group
python train.py --multirun training.lr=1e-3,3e-4,1e-4 # a 3-way grid search, one line

On top of that, every Hydra run gets its own timestamped output directory containing the fully resolved config it actually ran with -- overrides and all. So there's no more "wait, did I run this with lr=1e-3 or 3e-4?". The answer is sitting in a folder, frozen. Wire Hydra together with MLflow or wandb and you've got the full chain of custody: which config, which overrides, which code, which results. Nota bene: for me this is the single highest-leverage habit in the whole episode. Configs-as-files is boring, and boring is exactly what reproducibility is made of.

Docker: freezing the whole world

"But it works on my machine!" is the oldest lie in software, and ML tells it with extra conviction, because ML results are shockingly sensitive to the exact stack: a different CUDA version, a bumped NumPy, a subtly changed cuDNN, and your numbers wander. Docker ends the argument by shipping the entire environment -- OS libraries, CUDA, Python, packages, the lot -- as one image that runs identically anywhere.

# Dockerfile for a reproducible training environment
FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY src/ src/
COPY config/ config/

# Determinism knobs, baked into the image (more on these below)
ENV PYTHONHASHSEED=42
ENV CUBLAS_WORKSPACE_CONFIG=:4096:8

ENTRYPOINT ["python", "src/train.py"]

And the unglamorous hero of the whole file is this:

# requirements.txt -- pin EVERYTHING, exactly
torch==2.2.0
numpy==1.26.4
scikit-learn==1.4.0
mlflow==2.10.0
hydra-core==1.3.2
omegaconf==2.3.0

Look at those == signs and love them. numpy>=1.26 is not reproducible -- it means "whatever version happens to be newest the day someone rebuilds this", and a minor bump in a dependency can shift numerical behaviour just enough to move your accuracy by a percentage point. numpy==1.26.4 is a promise. Pin every direct dependency, and for real rigour pin the transitive ones too (with a tool like pip-tools or uv). A part from that, Docker also makes your work portable in the nicest way: the exact same image trains on your laptop, on a cloud GPU, and in CI, with byte-identical libraries every time.

Random seeds and the limits of determinism

Now the part everybody underestimates. Setting random.seed(42) feels like you've done the reproducibility thing, but it's barely the first of several seeds you need -- Python's random, NumPy, PyTorch on CPU, and PyTorch on CUDA each have their own generator:

import os
import random
import numpy as np
import torch

def set_seed(seed=42):
    """Pin every random source we can reach."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # Force deterministic kernels (slower, but bit-reproducible)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ["PYTHONHASHSEED"] = str(seed)
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    torch.use_deterministic_algorithms(True)

set_seed(42)

The cudnn.deterministic and use_deterministic_algorithms flags are doing something specific: they force PyTorch to pick deterministic implementations of operations that are, by default, non-deterministic for speed. Some GPU ops (think the atomicAdd inside a scatter) let many threads race to accumulate into the same memory, and floating-point addition is not associative -- (a + b) + c can differ from a + (b + c) in the last bits -- so the answer depends on which thread got there first. Forcing determinism serialises that and costs you speed. It's a real trade: bit-exact results, or a faster run. For a paper or a debugging session, take determinism; for a giant production sweep where you only care about the average, let it rip.

And here's the humbling bit, the thing I want you to walk away with: even with every seed pinned and every deterministic flag flipped, a model trained on an A100 may not reproduce bit-for-bit on a V100. Different GPU architectures order their floating-point operations differently, and those last-bit differences compound over millions of steps. That's not a bug in your code -- it's the hardware. True reproducibility has a hardware-shaped ceiling, and pretending otherwise is how people waste a weekend hunting a "bug" that is really just physics ;-)

Putting it all together

None of these tools is the answer on its own -- the reproducibility comes from stacking them. Here's the shape of a run that actually earns the word:

import subprocess
import mlflow
import hydra
from omegaconf import DictConfig, OmegaConf

def git_hash():
    return subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()[:8]

@hydra.main(config_path="config", config_name="train", version_base=None)
def main(cfg: DictConfig):
    set_seed(cfg.get("seed", 42))                       # same dice every time

    mlflow.set_experiment(cfg.experiment_name)
    with mlflow.start_run():
        # Log the FULLY RESOLVED config + the exact code version.
        mlflow.log_params(OmegaConf.to_container(cfg, resolve=True))
        mlflow.log_param("git_hash", git_hash())

        model = build_model(cfg.model)
        train(model, cfg.training)
        mlflow.log_metrics(evaluate(model, cfg.data.val_path))
        mlflow.pytorch.log_model(model, "model")

Read the log line for git_hash again, because it's the keystone. Same code (git hash) + same data (DVC version from #118) + same config (Hydra) + same environment (Docker) + same seed (set_seed) = same result. Knock out any one of those five and you're back to guessing. That's the entire discipline, compressed into an equation. Everything in this episode has been about making each of those five terms a thing you capture automatically instead of a thing you hope you remembered.

Let's recap

  • Experiment tracking (MLflow, wandb) auto-logs parameters, metrics and model artifacts so every run is searchable and comparable -- the cure for "which config was Tuesday's 94%?";
  • MLflow is open-source and self-hosted with zero fuss; wandb is a hosted service with slicker dashboards and built-in Bayesian sweeps -- pick one and use it from run number one;
  • Hydra turns scattered magic-number hyperparameters into composable YAML with command-line overrides and per-run output folders, so nobody ever asks "what did this run use?" again;
  • Docker freezes the entire software world (OS, CUDA, Python, packages) into one portable image -- and version-pin every dependency with ==, because >= is a reproducibility bug waiting to happen;
  • random seeds must be set for Python, NumPy, PyTorch AND CUDA together, and full determinism needs extra flags that trade speed for bit-exact results;
  • true reproducibility = same code + same data + same config + same environment + same seed -- and even then, different GPU hardware sets a last-bit ceiling you cannot code your way past.

Zoom out and notice the through-line of the last three episodes: #117 taught us to think in systems, #118 built the data machinery, and today we made the whole thing repeatable. What we've quietly assumed all along, though, is that once a model is trained we can just... run it. In the next stretch we'll poke hard at that assumption -- because a model that's correct but too slow, or too fat, or too hungry for hardware, is a model that never actually reaches a user. The road system is built; soon we start worrying about how fast the cars can actually go.

Bedankt en tot de volgende keer!

@scipio

stem stemsocial steemstem python programming
0
0
0.000
0 comments