path: root/experiments/run.py

"""
SDIL main training / diagnostics driver.

Trains one of {bp, dfa, sdil} (sdil with ablation flags) on MNIST/FashionMNIST,
logging the quantities that actually test the hypothesis:
  - train loss, test accuracy
  - per-hidden-layer cos(innovation r_l, -grad h_l)      <- the headline metric
  - cos(raw apical a_l, -grad)  and  cos(A_l c, -grad)   <- residualization ablation
  - single-step loss-decrease ratio vs exact GD

Everything is JSON-logged for later plotting.
"""
import argparse
import json
import os
import subprocess
import sys
import time

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.core import SDILNet, SDILConfig, sdil_step, neutral_p_update
from sdil.baselines import BPNet, dfa_config, evaluate
from sdil.local_baselines import FANet
from sdil import probes
from sdil.data import (get_dataset_splits, onehot, make_hierarchical,
                       make_teacher_student, make_tentmap,
                       split_training_loader)


REAL_DATASETS = ("mnist", "fmnist", "cifar10")
SYNTHETIC_DATASETS = ("teacher", "hierarchical", "tentmap")


def code_provenance():
    """Best-effort source revision metadata for reproducible result files."""
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    try:
        commit = subprocess.run(
            ["git", "rev-parse", "HEAD"], cwd=root, check=True,
            capture_output=True, text=True).stdout.strip()
        dirty = bool(subprocess.run(
            ["git", "status", "--porcelain", "--untracked-files=no"], cwd=root,
            check=True, capture_output=True, text=True).stdout.strip())
        return {"git_commit": commit, "git_dirty": dirty}
    except (OSError, subprocess.CalledProcessError):
        return {"git_commit": None, "git_dirty": None}


def device_sync(device):
    if str(device).startswith("cuda") and torch.cuda.is_available():
        torch.cuda.synchronize()


def reset_peak_memory(device):
    if str(device).startswith("cuda") and torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats(torch.device(device))


def hardware_report(device):
    report = {"device": str(device), "torch_version": torch.__version__}
    if str(device).startswith("cuda") and torch.cuda.is_available():
        cuda_device = torch.device(device)
        props = torch.cuda.get_device_properties(cuda_device)
        report.update({
            "cuda_device_name": props.name,
            "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
            "device_total_memory_bytes": props.total_memory,
            "peak_memory_allocated_bytes": torch.cuda.max_memory_allocated(cuda_device),
            "peak_memory_reserved_bytes": torch.cuda.max_memory_reserved(cuda_device),
        })
    else:
        report.update({
            "cuda_device_name": None,
            "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
            "device_total_memory_bytes": None,
            "peak_memory_allocated_bytes": None,
            "peak_memory_reserved_bytes": None,
        })
    return report


def calibration_work_per_event(net, cfg):
    """Hardware-independent causal-calibration work for one minibatch.

    Forward equivalents use affine multiply-add work as the denominator. The
    returned loss evaluations count calls producing per-example scalar losses;
    the caller multiplies by the actual minibatch size for scalar observations.
    """
    directions = cfg.pert_ndirs
    if cfg.pert_mode == "simultaneous":
        return {
            "batch_loss_evaluations": 2 * directions,
            # The implementation performs one duplicate clean forward plus one
            # batched +/- forward for every direction.
            "forward_equivalent_batches": 1.0 + 2.0 * directions,
            "perturbation_batch_expansion": 2 * directions,
        }

    full_work = sum(weight.numel() for weight in net.W)
    tail_work = 0
    for hidden_layer in range(net.L - 1):
        tail_work += sum(net.W[j].numel() for j in range(hidden_layer + 1, net.L))
    return {
        # One clean baseline loss plus +/- tail loss for every hidden layer.
        "batch_loss_evaluations": 1 + 2 * directions * (net.L - 1),
        "forward_equivalent_batches": 1.0 + 2.0 * directions * tail_work / full_work,
        "perturbation_batch_expansion": 1,
    }


def build(args, device):
    sizes = [args.n_in] + [args.width] * args.depth + [args.n_out]
    if args.mode == "bp":
        net = BPNet(sizes, act=args.act, device=device, seed=args.seed,
                    w_scale=args.w_scale, nuis_rho=0.0, residual=bool(args.residual),
                    predictor_mode=args.predictor_mode)
        cfg = SDILConfig(eta=args.eta, momentum=args.momentum)
        return net, cfg
    if args.mode == "fa":
        net = FANet(sizes, act=args.act, device=device, seed=args.seed,
                    w_scale=args.w_scale, nuis_rho=0.0,
                    residual=bool(args.residual),
                    predictor_mode=args.predictor_mode,
                    b_scale=args.feedback_scale)
        return net, SDILConfig(eta=args.eta, momentum=args.momentum)
    net = SDILNet(sizes, act=args.act, device=device, seed=args.seed,
                  w_scale=args.w_scale, a_scale=args.a_scale,
                  nuis_rho=args.nuis_rho, feedback=args.feedback,
                  residual=bool(args.residual), predictor_mode=args.predictor_mode,
                  traffic_mode=args.traffic_mode, nuis_seed=args.traffic_seed)
    if args.mode == "dfa":
        cfg = dfa_config(eta=args.eta, momentum=args.momentum)
    elif args.mode == "sdil":
        cfg = SDILConfig(
            eta=args.eta, eta_A=args.eta_A, eta_P=args.eta_P,
            use_residual=bool(args.use_residual), learn_A=bool(args.learn_A),
            learn_P=bool(args.learn_P), pert_sigma=args.pert_sigma,
            pert_every=args.pert_every, pert_ndirs=args.pert_ndirs,
            pert_mode=args.pert_mode,
            momentum=args.momentum, settle_steps=args.settle_steps,
            kappa=args.kappa, feedback=args.feedback,
            p_update_on_neutral=bool(args.p_neutral),
            normalize_delta=bool(args.normalize_delta),
            raw_scale_control=args.raw_scale_control)
    else:
        raise ValueError(args.mode)
    return net, cfg


def load_task(args, device):
    """Load a real dataset or construct one fixed synthetic task.

    ``task_seed`` controls examples and the target function; ``seed`` controls
    only the student initialization.  A depth sweep therefore compares models
    on exactly the same compositional problem instead of silently changing the
    teacher between model seeds.
    """
    if args.dataset in REAL_DATASETS:
        train, validation, test, n_in, n_out, split = get_dataset_splits(
            args.dataset, batch_size=args.batch_size, device=device,
            train_limit=args.train_examples or None,
            val_examples=args.val_examples, split_seed=args.split_seed)
        if args.eval_split == "validation":
            if validation is None:
                raise ValueError("--eval_split validation requires --val_examples > 0")
            evaluation = validation
        else:
            evaluation = test
        split["evaluation_split"] = args.eval_split
        return train, evaluation, n_in, n_out, split
    common = dict(
        n_train=args.task_train_examples,
        n_test=args.task_test_examples,
        seed=args.task_seed,
        batch_size=args.batch_size,
        device=device,
    )
    if args.dataset == "teacher":
        task = make_teacher_student(
            n_in=args.task_n_in, n_classes=args.task_classes,
            t_depth=args.teacher_depth, t_width=args.teacher_width,
            residual=bool(args.teacher_residual), **common)
    elif args.dataset == "hierarchical":
        task = make_hierarchical(
            levels=args.task_levels, n_classes=args.task_classes, **common)
    elif args.dataset == "tentmap":
        task = make_tentmap(
            levels=args.task_levels, n_in=args.task_n_in, **common)
    else:
        raise ValueError(f"unknown dataset: {args.dataset}")
    train, test, n_in, n_out = task
    validation = None
    split_details = {}
    if args.val_examples:
        train, validation, split_details = split_training_loader(
            train, args.val_examples, args.split_seed, args.batch_size)
    if args.eval_split == "validation":
        if validation is None:
            raise ValueError("--eval_split validation requires --val_examples > 0")
        evaluation = validation
    else:
        evaluation = test
    split = {
        "dataset": args.dataset,
        "task_seed": args.task_seed,
        "train_examples": len(train.x),
        "test_examples": args.task_test_examples,
        "evaluation_split": args.eval_split,
        "synthetic_generator": True,
    }
    split.update(split_details)
    return train, evaluation, n_in, n_out, split


def train(args):
    device = args.device
    torch.manual_seed(args.seed)
    train_loader, eval_loader, n_in, n_out, split = load_task(args, device)
    args.n_in, args.n_out = n_in, n_out
    net, cfg = build(args, device)
    # Reset after persistent data/model allocation so the peak includes the
    # resident training state plus transient perturbation-batch expansion.
    reset_peak_memory(device)

    # A fixed probe batch is loaded only when diagnostics are requested. In a
    # frozen test-only run with diagnostics disabled, the test set is therefore
    # not touched until the single final evaluation.
    px = py = poh = None
    if args.diagnostics != "none":
        px, py = next(iter(eval_loader))
        px, py = px[:args.probe_bs].to(device), py[:args.probe_bs].to(device)
        poh = onehot(py, n_out, device=device)

    log = {"args": vars(args), "split": split, "provenance": code_provenance(),
           "steps": [], "final": {}}
    step = 0
    prev_error = None
    device_sync(device)
    t0 = time.time()
    train_wall_s = 0.0
    eval_wall_s = 0.0
    diagnostics_wall_s = 0.0
    inline_diagnostics_wall_s = 0.0
    warmup_examples = 0
    ordinary_forward_examples = 0
    perturbation_events = 0
    calibration_batch_loss_evaluations = 0
    calibration_example_loss_evaluations = 0
    calibration_forward_equivalent_examples = 0.0
    perturbation_batch_expansion = 0

    # predictor warmup on neutral-period (c=0) drive, so P cancels the apical
    # nuisance before task plasticity relies on the residual (no-op when rho=0).
    if args.mode == "sdil" and args.learn_P and args.p_warmup_steps > 0 and args.nuis_rho > 0:
        device_sync(device)
        warmup_t0 = time.time()
        it = iter(train_loader)
        for _ in range(args.p_warmup_steps):
            try:
                wx, _ = next(it)
            except StopIteration:
                it = iter(train_loader)
                wx, _ = next(it)
            wx = wx.to(device)
            warmup_examples += wx.shape[0]
            neutral_p_update(net, wx, args.p_warmup_eta)
        device_sync(device)
        warmup_wall_s = time.time() - warmup_t0
    else:
        warmup_wall_s = 0.0
    for epoch in range(args.epochs):
        device_sync(device)
        train_t0 = time.time()
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            ordinary_forward_examples += x.shape[0]
            yoh = onehot(y, n_out, device=device)
            if args.mode == "bp":
                loss = net.bp_step(x, y, cfg.eta, momentum=cfg.momentum)
            elif args.mode == "fa":
                loss = net.fa_step(x, y, yoh, cfg.eta, momentum=cfg.momentum)
            else:
                loss, aux = sdil_step(net, x, y, yoh, cfg, step, prev_error=prev_error)
                prev_error = aux["error"]
                if args.mode == "sdil" and aux["did_pert"]:
                    event = calibration_work_per_event(net, cfg)
                    perturbation_events += 1
                    calibration_batch_loss_evaluations += event["batch_loss_evaluations"]
                    calibration_example_loss_evaluations += (
                        event["batch_loss_evaluations"] * x.shape[0])
                    calibration_forward_equivalent_examples += (
                        event["forward_equivalent_batches"] * x.shape[0])
                    perturbation_batch_expansion = max(
                        perturbation_batch_expansion,
                        event["perturbation_batch_expansion"])

            if step % args.log_every == 0:
                rec = {"step": step, "epoch": epoch, "train_loss": float(loss)}
                if args.diagnostics != "none":
                    device_sync(device)
                    diagnostics_t0 = time.time()
                if args.mode == "fa" and args.diagnostics != "none":
                    rec.update(probes.fa_alignment_report(net, px, py, poh))
                elif args.mode != "bp" and args.diagnostics != "none":
                    al = probes.alignment_report(net, px, py, poh, cfg)
                    rec["cos_r_negg"] = al["cos_r_negg"]
                    rec["cos_innovation_negg"] = al["cos_innovation_negg"]
                    rec["cos_apical_negg"] = al["cos_apical_negg"]
                    rec["cos_Ac_negg"] = al["cos_Ac_negg"]
                    rec["r_norm"] = al["r_norm"]
                    rec["traffic_norm"] = al["traffic_norm"]
                    rec["traffic_residual_norm"] = al["traffic_residual_norm"]
                    rec["traffic_r2"] = al["traffic_r2"]
                    if (args.mode == "sdil" and args.diagnostics == "full"
                            and step % (args.log_every * 5) == 0):
                        rec["ldr"] = probes.loss_decrease_ratio(net, px, py, poh, cfg, step)
                if args.diagnostics != "none":
                    device_sync(device)
                    elapsed = time.time() - diagnostics_t0
                    diagnostics_wall_s += elapsed
                    inline_diagnostics_wall_s += elapsed
                log["steps"].append(rec)
            step += 1
            if args.max_steps and step >= args.max_steps:
                break
        if args.max_steps and step >= args.max_steps:
            device_sync(device)
            train_wall_s += time.time() - train_t0
            break

        device_sync(device)
        train_wall_s += time.time() - train_t0

        should_evaluate = args.eval_every > 0 and (epoch + 1) % args.eval_every == 0
        record = {"epoch_end": epoch, "step": step}
        msg = f"[{args.tag}] epoch {epoch} step {step} loss {loss:.4f}"
        if should_evaluate:
            device_sync(device)
            eval_t0 = time.time()
            acc, tloss = evaluate(net, eval_loader)
            device_sync(device)
            eval_wall_s += time.time() - eval_t0
            metric = "val_acc" if args.eval_split == "validation" else "test_acc"
            msg += f" {metric} {acc:.4f}"
            record.update({"eval_split": args.eval_split,
                           "eval_acc": acc, "eval_loss": tloss})
            if args.eval_split == "test":
                record.update({"test_acc": acc, "test_loss": tloss})
            else:
                record.update({"val_acc": acc, "val_loss": tloss})
        if args.diagnostics != "none":
            device_sync(device)
            diagnostics_t0 = time.time()
            if args.mode == "fa":
                al = probes.fa_alignment_report(net, px, py, poh)
                meancos = sum(al["cos_fa_negg"]) / len(al["cos_fa_negg"])
                msg += f" mean_cos(fa,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_fa_negg']]}"
            elif args.mode != "bp":
                al = probes.alignment_report(net, px, py, poh, cfg)
                meancos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"])
                msg += f" mean_cos(r,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_r_negg']]}"
            device_sync(device)
            diagnostics_wall_s += time.time() - diagnostics_t0
        print(msg, flush=True)
        log["steps"].append(record)

    device_sync(device)
    eval_t0 = time.time()
    acc, tloss = evaluate(net, eval_loader)
    device_sync(device)
    eval_wall_s += time.time() - eval_t0
    log["final"] = {"eval_split": args.eval_split, "eval_acc": acc, "eval_loss": tloss,
                    "wall_s": None}
    if args.eval_split == "test":
        log["final"].update({"test_acc": acc, "test_loss": tloss})
    else:
        log["final"].update({"val_acc": acc, "val_loss": tloss})
    if args.mode == "fa" and args.diagnostics != "none":
        device_sync(device)
        diagnostics_t0 = time.time()
        log["final"].update(probes.fa_alignment_report(net, px, py, poh))
        device_sync(device)
        diagnostics_wall_s += time.time() - diagnostics_t0
    elif args.mode != "bp" and args.diagnostics != "none":
        device_sync(device)
        diagnostics_t0 = time.time()
        al = probes.alignment_report(net, px, py, poh, cfg)
        log["final"]["cos_r_negg"] = al["cos_r_negg"]
        log["final"]["cos_innovation_negg"] = al["cos_innovation_negg"]
        log["final"]["cos_apical_negg"] = al["cos_apical_negg"]
        log["final"]["cos_Ac_negg"] = al["cos_Ac_negg"]
        log["final"]["traffic_norm"] = al["traffic_norm"]
        log["final"]["traffic_residual_norm"] = al["traffic_residual_norm"]
        log["final"]["traffic_r2"] = al["traffic_r2"]
        device_sync(device)
        diagnostics_wall_s += time.time() - diagnostics_t0
    device_sync(device)
    log["final"]["wall_s"] = time.time() - t0
    log["timing"] = {
        "warmup_wall_s": warmup_wall_s,
        "training_loop_wall_s": train_wall_s,
        "diagnostics_wall_s": diagnostics_wall_s,
        "inline_diagnostics_wall_s": inline_diagnostics_wall_s,
        "optimizer_wall_s_excluding_inline_diagnostics": max(
            0.0, train_wall_s - inline_diagnostics_wall_s),
        "evaluation_wall_s": eval_wall_s,
    }
    log["cost"] = {
        "train_steps": step,
        "ordinary_training_forward_examples": ordinary_forward_examples,
        "predictor_warmup_forward_examples": warmup_examples,
        "perturbation_events": perturbation_events,
        "calibration_batch_loss_evaluations": calibration_batch_loss_evaluations,
        "calibration_example_loss_evaluations": calibration_example_loss_evaluations,
        "calibration_forward_equivalent_examples": calibration_forward_equivalent_examples,
        "training_forward_equivalent_examples": (
            ordinary_forward_examples + warmup_examples
            + calibration_forward_equivalent_examples),
        "max_perturbation_batch_expansion": perturbation_batch_expansion,
    }
    log["hardware"] = hardware_report(device)
    os.makedirs(args.outdir, exist_ok=True)
    outpath = os.path.join(args.outdir, f"{args.tag}.json")
    with open(outpath, "w") as f:
        json.dump(log, f)
    print(f"[{args.tag}] DONE {args.eval_split}_acc={acc:.4f} -> {outpath}", flush=True)
    return log


def get_args():
    p = argparse.ArgumentParser()
    p.add_argument("--mode", default="sdil", choices=["bp", "fa", "dfa", "sdil"])
    p.add_argument("--dataset", default="mnist",
                   choices=list(REAL_DATASETS + SYNTHETIC_DATASETS))
    p.add_argument("--depth", type=int, default=3)          # hidden layers
    p.add_argument("--width", type=int, default=256)
    p.add_argument("--act", default="tanh", choices=["tanh", "gelu", "silu", "relu"])
    p.add_argument("--residual", type=int, default=0)       # skip connections (deep no-BN)
    p.add_argument("--epochs", type=int, default=15)
    p.add_argument("--batch_size", type=int, default=128)
    p.add_argument("--train_examples", type=int, default=0,
                   help="0 uses the full training split")
    p.add_argument("--val_examples", type=int, default=0,
                   help="stratified validation examples held out from training")
    p.add_argument("--split_seed", type=int, default=2027)
    p.add_argument("--eval_split", default="test", choices=["validation", "test"])
    p.add_argument("--task_seed", type=int, default=0,
                   help="fixed target/data seed, separate from student --seed")
    p.add_argument("--task_train_examples", type=int, default=50000)
    p.add_argument("--task_test_examples", type=int, default=10000)
    p.add_argument("--task_levels", type=int, default=8)
    p.add_argument("--task_n_in", type=int, default=128)
    p.add_argument("--task_classes", type=int, default=10)
    p.add_argument("--teacher_depth", type=int, default=8)
    p.add_argument("--teacher_width", type=int, default=64)
    p.add_argument("--teacher_residual", type=int, default=1)
    p.add_argument("--eta", type=float, default=0.05)
    p.add_argument("--eta_A", type=float, default=0.02)
    p.add_argument("--eta_P", type=float, default=0.002)
    p.add_argument("--momentum", type=float, default=0.9)
    p.add_argument("--w_scale", type=float, default=1.0)
    p.add_argument("--a_scale", type=float, default=1.0)
    p.add_argument("--feedback_scale", type=float, default=1.0)
    p.add_argument("--pert_sigma", type=float, default=1e-2)
    p.add_argument("--pert_every", type=int, default=4)
    p.add_argument("--pert_ndirs", type=int, default=4)
    p.add_argument("--pert_mode", default="layerwise", choices=["layerwise", "simultaneous"])
    p.add_argument("--use_residual", type=int, default=1)
    p.add_argument("--raw_scale_control", default="none",
                   choices=["none", "match_innovation_norm"])
    p.add_argument("--learn_A", type=int, default=1)
    p.add_argument("--learn_P", type=int, default=1)
    p.add_argument("--p_neutral", type=int, default=1)      # P update on neutral (c=0) drive
    p.add_argument("--p_warmup_steps", type=int, default=200)  # pre-task neutral P warmup
    p.add_argument("--p_warmup_eta", type=float, default=0.05)
    p.add_argument("--nuis_rho", type=float, default=0.0)
    p.add_argument("--traffic_seed", type=int, default=1234,
                   help="ordinary apical-traffic projection seed; independent of model/data seeds")
    p.add_argument("--traffic_mode", default="soma",
                   choices=["none", "soma", "topdown", "mixed"])
    p.add_argument("--predictor_mode", default="diagonal", choices=["diagonal", "full"])
    p.add_argument("--normalize_delta", type=int, default=0)
    p.add_argument("--settle_steps", type=int, default=0)
    p.add_argument("--kappa", type=float, default=0.0)
    p.add_argument("--feedback", default="error", choices=["error", "error_deriv"])
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--log_every", type=int, default=50)
    p.add_argument("--diagnostics", default="full", choices=["none", "alignment", "full"])
    p.add_argument("--eval_every", type=int, default=1,
                   help="0 evaluates only once at the end")
    p.add_argument("--max_steps", type=int, default=0)      # 0 = no cap (smoke only)
    p.add_argument("--probe_bs", type=int, default=512)
    p.add_argument("--outdir", default="results")
    p.add_argument("--tag", default="sdil_run")
    return p.parse_args()


if __name__ == "__main__":
    train(get_args())