Hyperparameter Search and Experiment Tracking

Hyperparameter search is the disciplined process of choosing values that control training but are not learned by gradient descent. Experiment tracking is the record that lets you explain which values, code, data split, metric, and artifact produced a result. In this PyTorch course, the outcome is practical: improve a model without accidentally training on the validation set, losing the winning configuration, or comparing runs that were not actually comparable.

Purpose and Outcome

A PyTorch model learns parameters such as weights and biases from batches. Hyperparameters sit outside that learned state: learning rate, optimizer choice, batch size, weight decay, dropout, scheduler settings, architecture width, data augmentation strength, number of epochs, early-stopping patience, and loss-function options. Search asks, “Which configuration should I train next?” Tracking asks, “What exactly happened when I trained it?” Together they turn model improvement from memory and screenshots into reproducible evidence.

The key boundary is the validation set. You may use validation metrics to select hyperparameters, but the final test set must remain untouched until the end. If you repeatedly use the test set during search, you are fitting decisions to test examples even though no tensor gradient flows through them. The reported performance then becomes optimistic.

How Search Works Internally

A search system repeatedly creates a trial. Each trial contains a configuration, initializes training, records metrics, and ends with a state such as completed, pruned, failed, or timed out. The search algorithm decides the next configuration from a search space. The trainer consumes the configuration and emits measurements. The tracker persists the record so later ranking does not depend on console scrollback.

Grid search enumerates every combination from finite lists. It is easy to audit and parallelize, but grows multiplicatively: three learning rates, two batch sizes, and two decay values already make twelve trials. Random search samples from distributions. It often beats grid search when only a few hyperparameters matter because it explores more distinct values for each dimension. Bayesian and bandit methods use previous results to choose promising regions or stop weak trials early. Those methods can be efficient, but they also introduce more machinery: samplers, priors, warmup rules, pruning thresholds, and a higher risk that noisy validation metrics mislead the scheduler.

In PyTorch itself, there is no special hyperparameter object. A trial is ordinary Python data passed into ordinary constructors: torch.optim.AdamW(..., lr=...), DataLoader(..., batch_size=...), nn.Dropout(p=...), and training-loop conditions such as epoch count or patience. That is powerful because every training choice is programmable, but it means the engineer must make run identity, seeding, logging, and artifact storage explicit.

API Anatomy

A reliable trial function usually has five parts. First, it accepts a plain configuration dictionary or typed object. Second, it sets seeds and creates train, validation, and test loaders without leakage. Third, it builds a fresh model and optimizer for that configuration. Fourth, it logs scalar metrics such as train loss, validation loss, accuracy, F1 score, learning rate, duration, and device information. Fifth, it returns the selection metric and stores the model artifact only when needed.

Common tracking fields are run_id, parent_run_id, config, git_commit, data_version, seed, epoch_metrics, final_metrics, artifact_uri, status, and notes. The important design point is that metrics and configuration are separate. A learning rate is an input. A validation loss is an output. Mixing them makes filtering and comparison brittle.

Example 1: Exhaustive Grid Enumeration

This example builds the complete list of grid-search trials. It does not train a network yet; it proves the search space is what you think it is before you spend GPU time.

from itertools import product

search_space = {
    "lr": [0.1, 0.01, 0.001],
    "batch_size": [32, 64],
    "weight_decay": [0.0, 0.0001],
}

keys = list(search_space)
configs = [dict(zip(keys, values)) for values in product(*(search_space[k] for k in keys))]

print(len(configs))
print(configs[0])
print(configs[-1])

The deterministic output starts with 12, then prints the first and last configuration. The behavior shows why grid search is transparent: there are exactly twelve planned trials, and every value is visible before training starts. It also shows the trade-off: adding one more four-value hyperparameter would make forty-eight trials.

Example 2: Random Search Over Continuous Values

Random search is useful when a small list is too coarse. Learning rate is often better searched on a logarithmic scale because 0.0001 to 0.001 is as meaningful as 0.01 to 0.1. The scoring function below is a deterministic stand-in for validation loss, so the example can show exact behavior without needing a GPU.

import random

random.seed(7)
space = [
    {"lr": 10 ** random.uniform(-4, -1), "dropout": random.uniform(0.0, 0.5)}
    for _ in range(4)
]

def score(config: dict[str, float]) -> float:
    lr_penalty = abs(config["lr"] - 0.01) * 30
    dropout_penalty = abs(config["dropout"] - 0.2)
    return round(0.35 + lr_penalty + dropout_penalty, 4)

ranked = sorted((score(cfg), cfg) for cfg in space)
print(ranked[0][0])
print(round(ranked[0][1]["lr"], 5), round(ranked[0][1]["dropout"], 3))

With the fixed seed, this prints a best score of 0.5447 and a rounded configuration near 0.00897 learning rate and 0.036 dropout. In a real PyTorch loop, score would be replaced by training and validation. The search idea is unchanged: sample candidates, evaluate them consistently, and rank by a metric chosen before looking at results.

Example 3: Minimal Experiment Tracking

A tracker can be a hosted tool, a database, a local file, or the in-memory buffer shown while learning the mechanism. JSON Lines is convenient because each trial is one appendable record. This example records three finished trials and reloads the best one.

import io
import json

trials = [
    {"run_id": "trial-001", "lr": 0.1, "batch_size": 32, "val_loss": 0.82},
    {"run_id": "trial-002", "lr": 0.01, "batch_size": 32, "val_loss": 0.41},
    {"run_id": "trial-003", "lr": 0.001, "batch_size": 64, "val_loss": 0.55},
]

buffer = io.StringIO()
for trial in trials:
    buffer.write(json.dumps(trial, sort_keys=True) + "\n")

loaded = [json.loads(line) for line in buffer.getvalue().splitlines()]
best = min(loaded, key=lambda row: row["val_loss"])
print(best["run_id"], best["lr"], best["val_loss"])

The output is trial-002 0.01 0.41. The important behavior is not the storage backend; it is that the comparison can be reproduced from structured records. Real trackers add dashboards, artifacts, tags, nested runs, and concurrent writers, but the same record structure remains: immutable trial inputs plus measured outputs.

Example 4: A PyTorch Trial Function

The next example trains a tiny regression model. Each configuration builds a fresh model and optimizer. Reusing the same model across trials would contaminate the search because the second trial would start from parameters already improved by the first.

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def make_data(n: int = 128) -> TensorDataset:
    torch.manual_seed(0)
    x = torch.linspace(-2, 2, n).unsqueeze(1)
    y = 3 * x - 0.5 + 0.2 * torch.randn_like(x)
    return TensorDataset(x, y)

def train_trial(config: dict[str, float | int]) -> dict[str, float]:
    torch.manual_seed(123)
    loader = DataLoader(make_data(), batch_size=int(config["batch_size"]), shuffle=True)
    model = nn.Sequential(nn.Linear(1, 8), nn.ReLU(), nn.Linear(8, 1))
    opt = torch.optim.AdamW(model.parameters(), lr=float(config["lr"]), weight_decay=float(config["weight_decay"]))
    loss_fn = nn.MSELoss()
    for _ in range(25):
        for xb, yb in loader:
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()
    with torch.no_grad():
        x_val = torch.tensor([[-1.5], [0.0], [1.5]])
        y_val = 3 * x_val - 0.5
        val_loss = loss_fn(model(x_val), y_val).item()
    return {"val_loss": round(val_loss, 4)}

for cfg in [
    {"lr": 0.01, "batch_size": 16, "weight_decay": 0.0},
    {"lr": 0.001, "batch_size": 32, "weight_decay": 0.0001},
]:
    result = train_trial(cfg)
    print(cfg, result)

The printed losses can vary slightly across platforms, but both lines should print a configuration dictionary and a val_loss dictionary. The pattern scales to image or text models: isolate one trial, pass hyperparameters into constructors, validate on held-out data, and return the metric used for ranking.

Design Choices and Trade-offs

Choose the search budget first. A large model with long epochs may need a coarse random search, early stopping, and a smaller proxy dataset. A small model may justify grid search because auditability is worth more than sampler sophistication. Avoid comparing a two-hour full-data run against a five-minute proxy run unless the tracker labels those regimes clearly.

Pick one primary selection metric before searching. Accuracy can hide minority-class failure; validation loss can favor calibrated probabilities even when top-line accuracy is unchanged; F1 can be unstable on small validation sets. Secondary metrics should be logged, but the winning rule should not move after seeing the leaderboard.

Decide whether trials share anything. Sharing pretrained weights can reduce cost, but it changes the meaning of the search because initialization is no longer independent. Parallel trials improve throughput, but they require unique run directories and safe artifact names. Early pruning saves compute, but may kill configurations that learn slowly and finish strong.

Failure Modes and Troubleshooting

Symptom: validation improves during search, but test performance is disappointing. Cause: the validation set has become an optimization target after many trials, or preprocessing was fit on all data. Diagnose: inspect split creation, count how often validation was queried, and check whether scalers, tokenizers, or augmentations were fit only on training data. Correct: rebuild leakage-safe splits, keep a final untouched test set, and consider nested cross-validation for small datasets.

Symptom: two runs with the same configuration disagree wildly. Cause: uncontrolled seeds, nondeterministic kernels, different data order, or a changed dataset. Diagnose: compare logged seed, data version, package environment, device, and sampler settings. Correct: log those fields, set seeds at trial start, and report mean and variance across repeated trials when noise is material.

Symptom: the best run cannot be loaded later. Cause: metrics were tracked but model state, preprocessing vocabulary, or architecture code was not versioned. Diagnose: follow the recorded artifact URI and attempt inference from a clean process. Correct: store state_dict, configuration, preprocessing objects, and code version together, and verify reload as part of the lab.

Symptom: search crashes after several hours and restarts from trial one. Cause: trial state is only in memory. Diagnose: check whether completed run IDs are persisted before the next trial begins. Correct: append trial records as they finish and make the launcher skip completed IDs on resume.

Security, Performance, and Reliability

Experiment logs often contain dataset paths, labels, example identifiers, prompts, or user-derived metadata. Log enough to reproduce the run, but avoid raw sensitive records. If using a remote tracker, treat API tokens as credentials, keep them outside notebooks, and avoid uploading private artifacts by default.

Performance depends on scheduling as much as math. Data loading can dominate short trials, so tune num_workers and cache immutable preprocessing when appropriate. GPU memory failures often come from trying larger batch sizes during search; catch those failures, mark the trial failed, and continue instead of losing the whole sweep. Reliability improves when every trial has its own directory, atomic metric writes, and a status transition from running to completed or failed.

Hands-on Lab

Prerequisites: Python with PyTorch installed, a terminal, and a writable scratch directory. Use CPU if no GPU is available; the lab is intentionally small.

  1. Create a new script named search_lab.py and paste the PyTorch trial function from Example 4.
  2. Add a list of four configurations that vary lr, batch_size, and weight_decay.
  3. For each configuration, create a unique run_id, call train_trial, and append one JSON Lines record containing run_id, config, and final_metrics.
  4. After all trials finish, reload the JSON Lines file, select the record with the smallest final_metrics.val_loss, and print its run ID and configuration.
  5. Verification: rerun the script from a clean terminal. Confirm that every trial has one record, every record has the same keys, and the selected best run exists in the file. Also intentionally set one learning rate to a string and confirm the script fails before writing a misleading completed record.
  6. Cleanup: delete the scratch run directory or archive it with the date, dataset description, and code version if you want to keep the evidence.

Assessment Exercises

  1. You have a validation leaderboard after 200 random trials and one untouched test set. What exact steps would you take before reporting final performance, and why?
  2. A teammate wants to tune learning rate, batch size, optimizer, dropout, and model width with grid search. Estimate how trial count grows and propose a cheaper search plan.
  3. Design a tracking record for an image classifier. Which fields are inputs, which are outputs, and which are artifacts?
  4. A pruner stops all low-learning-rate trials after one epoch. Explain when that is efficient and when it biases the search against good configurations.
  5. Given two runs with equal validation accuracy but different validation loss and inference latency, define a defensible selection rule for a deployment target.

Summary

Hyperparameter search changes PyTorch training by wrapping the training loop in repeated, comparable trials. Experiment tracking makes those trials durable: configurations, metrics, seeds, data versions, code versions, and artifacts survive beyond the notebook session. Use grid search when auditability and small spaces matter, random or adaptive search when continuous spaces and budget matter, and always protect the final test set from search decisions.