scripts/claim6_scale_probe.py

"""Claim 6 scale probe: only the decisive stepsizes (1e-4 for both methods, plus the
baseline at 1e-5), at 25000 training images x 4 epochs = 100k batch-1 SGD steps.

Purpose: check whether the baseline's floating-point exception rate at eta = 1e-4 keeps
rising with the number of steps (1/3 at 50k steps), i.e. whether the paper's "baseline
requires eta = 1e-5" is a pure horizon effect.
"""

import json
import os
import sys
import time
from multiprocessing import Pool

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import claim6_uot_dro as m

OUT = m.OUT
TAG = "_xl"
N_TRAIN, EPOCHS, SEEDS = 25000, 4, 3

_CFG = {}


def _init():
    _CFG["data"] = m.load_data(N_TRAIN, 1000, 2000, 2000)
    _CFG["tag"] = TAG


def _worker(t):
    return m.train(*t, cfg=_CFG)


if __name__ == "__main__":
    os.makedirs(os.path.join(OUT, "uot_traj" + TAG), exist_ok=True)
    tasks = []
    for s in range(SEEDS):
        tasks.append(("baseline", 1e-4, s, None, N_TRAIN, EPOCHS))
        tasks.append(("baseline", 1e-5, s, None, N_TRAIN, EPOCHS))
        tasks.append(("proposed", 1e-4, s, 0.1, N_TRAIN, EPOCHS))
    t0 = time.time()
    with Pool(len(tasks), initializer=_init) as pool:
        recs = pool.map(_worker, tasks, chunksize=1)
    d = m.load_data(N_TRAIN, 1000, 2000, 2000)
    with open(os.path.join(OUT, f"claim6_uot_dro{TAG}.json"), "w") as fh:
        json.dump(
            {
                "config": {"n_train": N_TRAIN, "epochs": EPOCHS, "seeds": SEEDS},
                "noise_rate": d["noise_rate"],
                "elapsed": time.time() - t0,
                "runs": recs,
            },
            fh,
        )
    for meth, lr in [("baseline", 1e-4), ("baseline", 1e-5), ("proposed", 1e-4)]:
        rs = [r for r in recs if r["method"] == meth and r["lr"] == lr]
        fin = [
            r["final_F"] for r in rs if not r["diverged"] and r["final_F"] is not None
        ]
        print(
            meth,
            lr,
            "fpe",
            f"{sum(r['diverged'] for r in rs)}/{len(rs)}",
            [r["fpe_iter"] for r in rs],
            "F",
            None if not fin else round(float(np.mean(fin)), 4),
            flush=True,
        )