# Experiments for "Debugging a ResNet-50 Fine-Tuning Script": the notebook's code as written, then one fix at a time.
# Run in the task folder (DummyData/, imagenet_classes.txt, pretrained_models/resnet50-19c8e357.pth); PyTorch 2.14, CPU, seeds 0-2.
#   python debug_finetune.py OUT.json
import sys, json, copy, warnings, torch, torch.nn as nn, torch.optim as optim, torchvision.transforms as T
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from torchvision.models import resnet50
warnings.filterwarnings("ignore", category=UserWarning)
names = [l.split(",")[1].strip() for l in open("imagenet_classes.txt") if l.strip()]
IMNET = {0: 282, 1: 953, 2: 100}                                       # the notebook's get_imagenet_class
NB_NORM, IN_NORM = ([0.5] * 3, [0.1] * 3), ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
def data(norm):
    tf = T.Compose([T.Resize((224, 224)), T.ToTensor(), T.Normalize(*norm)])
    return ImageFolder("DummyData/train", tf), ImageFolder("DummyData/test", tf)
base = resnet50()
base.load_state_dict(torch.load("pretrained_models/resnet50-19c8e357.pth", map_location="cpu", weights_only=False))   # official file, legacy format
def evaluate(m, ds):
    m.eval(); rows = []
    with torch.no_grad():
        for x, y in DataLoader(ds, batch_size=1):
            p = torch.softmax(m(x), 1)[0]; top = p.topk(3); want = IMNET[y.item()]
            rows.append({"true": names[want], "top": [[names[i], round(v, 4)] for v, i in zip(top.values.tolist(), top.indices.tolist())], "p_true": round(p[want].item(), 4), "ok": top.indices[0].item() == want})
    return rows
def run(cfg, seed):
    torch.manual_seed(seed); tr, te = data(cfg["norm"]); m = copy.deepcopy(base)
    if cfg.get("head"): m.fc = nn.Linear(2048, 3)
    if cfg.get("freeze"):
        for n, p in m.named_parameters(): p.requires_grad = n.startswith("fc.")
    opt = optim.SGD([p for p in m.parameters() if p.requires_grad], lr=cfg["lr"], momentum=cfg["mom"])
    crit = nn.MSELoss() if cfg["loss"] == "mse" else nn.CrossEntropyLoss(); hist = []
    for ep in range(cfg.get("epochs", 3)):
        m.train()
        if cfg.get("bn_eval"): [mod.eval() for mod in m.modules() if isinstance(mod, nn.BatchNorm2d)]
        tot = 0.0
        for x, y in DataLoader(tr, batch_size=1, shuffle=True):
            t = torch.tensor([IMNET[v] for v in y.tolist()]) if cfg.get("map") else y
            if cfg.get("zero"): opt.zero_grad()
            loss = crit(m(x), t); loss.backward(); opt.step(); tot += loss.item()
        hist.append(tot / len(tr))
    if cfg.get("head"):
        m.eval(); rows = []
        with torch.no_grad():
            for x, y in DataLoader(te, batch_size=1):
                p = torch.softmax(m(x), 1)[0]; rows.append({"true": te.classes[y.item()], "top": [[te.classes[i], round(p[i].item(), 4)] for i in p.argsort(descending=True).tolist()], "p_true": round(p[y.item()].item(), 4), "ok": p.argmax().item() == y.item()})
        return hist, rows
    return hist, evaluate(m, te)
NB = dict(norm=NB_NORM, loss="mse", lr=0.9, mom=0.001)
CFG = [
    ("A", "The notebook as written (on PyTorch 2.14, once the weights load)", dict(NB)),
    ("B", "+ cross-entropy, labels mapped to ImageNet ids", dict(NB, loss="ce", map=True)),
    ("C", "+ zero_grad, + ImageNet normalisation (still lr 0.9, momentum 0.001)", dict(NB, loss="ce", map=True, zero=True, norm=IN_NORM)),
    ("D", "+ lr 0.001, momentum 0.9", dict(loss="ce", map=True, zero=True, norm=IN_NORM, lr=0.001, mom=0.9)),
    ("E", "D with BatchNorm frozen (eval mode)", dict(loss="ce", map=True, zero=True, norm=IN_NORM, lr=0.001, mom=0.9, bn_eval=True)),
    ("F", "E with lr 0.0001", dict(loss="ce", map=True, zero=True, norm=IN_NORM, lr=0.0001, mom=0.9, bn_eval=True)),
    ("G", "Frozen backbone (requires_grad=False), new 3-class head, lr 0.01, 10 epochs", dict(loss="ce", map=False, zero=True, norm=IN_NORM, lr=0.01, mom=0.9, head=True, freeze=True, epochs=10)),
    ("H", "G with BatchNorm also frozen (eval mode)", dict(loss="ce", map=False, zero=True, norm=IN_NORM, lr=0.01, mom=0.9, head=True, freeze=True, bn_eval=True, epochs=10)),
]
out = {"pretrained": {}, "configs": []}
tr_nb, te_nb = data(NB_NORM); tr_in, te_in = data(IN_NORM)
out["pretrained"]["notebook_norm"] = evaluate(base, te_nb); out["pretrained"]["imagenet_norm"] = evaluate(base, te_in)
x_nb = torch.stack([te_nb[i][0] for i in range(3)]); x_in = torch.stack([te_in[i][0] for i in range(3)])
out["inputs"] = {"notebook": [round(x_nb.min().item(), 3), round(x_nb.max().item(), 3), round(x_nb.std().item(), 3)], "imagenet": [round(x_in.min().item(), 3), round(x_in.max().item(), 3), round(x_in.std().item(), 3)]}
print("classes", tr_nb.class_to_idx, "| input range notebook", out["inputs"]["notebook"], "imagenet", out["inputs"]["imagenet"])
for k, rows in out["pretrained"].items(): print(f"pretrained, no training, {k}:", [(r["true"], r["top"][0], r["ok"]) for r in rows])
# what MSELoss does to the logits: it pulls all 1,000 of them toward the class index
m = copy.deepcopy(base); m.train(); x, y = tr_nb[0]; z = m(x[None]).detach()
print(f"MSE target broadcast: logits mean {z.mean():.3f}, sd {z.std():.3f}; loss vs target 0 = mean (z - 0)^2 = {((z - 0) ** 2).mean():.3f} (nn.MSELoss gives {nn.MSELoss()(z, torch.tensor([0])):.3f})")
for key, desc, cfg in CFG:
    res = [run(cfg, s) for s in range(3)]
    ok = [sum(r["ok"] for r in rows) for _, rows in res]
    out["configs"].append({"key": key, "desc": desc, "loss": [[round(v, 4) if v == v else None for v in h] for h, _ in res], "test": res[0][1], "ok_per_seed": ok})
    print(f"{key}. {desc}\n   loss per epoch (seed 0): {[round(v, 4) for v in res[0][0]]}; correct of 3 per seed: {ok}\n   seed 0:", "; ".join(f"{r['true']} -> {r['top'][0][0]} {r['top'][0][1]:.2f}" for r in res[0][1]))
def clean(o):   # NaN (from the diverged runs) is not valid JSON: write null
    if isinstance(o, float) and o != o: return None
    if isinstance(o, dict): return {k: clean(v) for k, v in o.items()}
    if isinstance(o, list): return [clean(v) for v in o]
    return o
if len(sys.argv) > 1: json.dump(clean(out), open(sys.argv[1], "w"), separators=(",", ":"))