Debug a PyTorch Training Run: NaN Loss, Broadcasting Bugs, Learning Rate and Data Leaks

"""Your debugging tools. Each one answers a question about the training run in train.py."""
import math

import numpy as np
import torch

import train as T


# ---------- Step 1: exploding loss ----------
def grad_norm(model):
    """The size of the whole gradient: the square root of the sum of every parameter gradient squared (the L2
    norm over all parameters together). Parameters without a gradient are skipped."""
    # TODO (Step 1): add up p.grad squared over every parameter that has a gradient; return the square root.
    raise NotImplementedError("Step 1: write grad_norm()")


# ---------- Step 2: a model that learns nothing ----------
def baseline_loss(y):
    """The loss of the laziest model: predict the mean of y for every row. MSE of that = the variance of y."""
    raise NotImplementedError("baseline_loss() arrives in Step 2")


def overfit_one_batch(xb, yb, steps=300, lr=1e-2, seed=0):
    """Train a fresh T.make_model() on this one batch only, with Adam(lr) and T.compute_loss, for `steps` steps.
    Returns the final loss. A working pipeline can memorise one batch almost perfectly."""
    raise NotImplementedError("overfit_one_batch() arrives in Step 2")


# ---------- Step 3: the learning rate ----------
def lr_range_test(X, y, lrs, steps=200, seed=0):
    """For every learning rate: torch.manual_seed(seed), a fresh T.make_model(), and T.train() it until `steps`
    batches have run (T.train's on_batch sees each loss). Score = the mean of the last 50 batch losses, or
    float("inf") if any of them is not finite. Returns {lr: score}."""
    raise NotImplementedError("lr_range_test() arrives in Step 3")


def pick_lr(results):
    """The learning rate one step below the one with the lowest score (in the sorted list of tried rates): the
    best rate in a short test sits close to the edge where training becomes unstable."""
    raise NotImplementedError("pick_lr() arrives in Step 3")


# ---------- Step 4: too good to be true ----------
def single_feature_r2(df, features):
    """How much of the target each column explains on its own: the squared correlation between the column and
    T.TARGET, over the rows of df. Returns {feature: r2}."""
    raise NotImplementedError("single_feature_r2() arrives in Step 4")


def suspicious(scores, threshold=0.95):
    """The features that explain more than `threshold` of the target on their own, best first."""
    raise NotImplementedError("suspicious() arrives in Step 4")


# ---------- Step 5: a preflight that catches all four ----------
def preflight(df):
    """Run before every training job, on the meter data df, using train.py as it is now (T.FEATURES, T.CONFIG,
    T.Prep, T.compute_loss, T.make_model). Returns a list of (check, message) for every problem found; empty
    when the job is safe to run. The checks:
      "scaling":       after T.Prep, every input column and the target have |mean| < 0.1 and 0.5 < std < 2.
      "overfit":       overfit_one_batch() on the first 32 training rows ends below 0.05.
      "learning_rate": T.CONFIG["lr"] is higher than the best rate of lr_range_test() over LR_GRID.
      "leak":          suspicious(single_feature_r2()) on the training rows finds any feature.
    """
    raise NotImplementedError("preflight() arrives in Step 5")


LR_GRID = [1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 0.1, 0.3, 1.0]