"""Train PathoCNN from scratch on PathMNIST — the exact script that produced the released weights. Protocol: - PathMNIST official splits (train 89,996 / val 10,004 / test 7,180), downloaded from Zenodo (record 10519652, md5 a8b06965200029087d5bd730944a56c1). - Inputs: /255, normalized per channel with the train-split mean/std (the rounded contract values in pathocnn.infer.TRAIN_MEAN/TRAIN_STD). - Augmentation (train only): random dihedral transforms — 90-degree rotations and flips. Histology patches have no canonical orientation, so the full D4 group is label-preserving for this data. No color jitter: H&E stain color carries class signal. - AdamW (lr 2e-3, weight decay 1e-4, cosine schedule), batch 256, 12 epochs, plain cross-entropy (train split max/min class ratio is 1.63, so no class weighting), seed 42 throughout. - Checkpoint selected by best validation accuracy; the test split is evaluated exactly once, after selection, and recorded in the metrics JSON. examples/verify_metrics.py later reproduces those test numbers from the shipped weights. Usage: python examples/train.py --data /path/to/pathmnist.npz --out /path/to/workdir If the npz is missing, download it first, e.g.: curl -L -o pathmnist.npz "https://zenodo.org/records/10519652/files/pathmnist.npz?download=1" (Python HTTP clients can stall on Zenodo through some proxies; curl is reliable.) """ import argparse import json import os import random import sys import time import numpy as np sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) import torch # noqa: E402 import torch.nn as nn # noqa: E402 from sklearn.metrics import (accuracy_score, balanced_accuracy_score, # noqa: E402 confusion_matrix, precision_score, recall_score, roc_auc_score) from pathocnn.infer import LABELS, TRAIN_MEAN, TRAIN_STD # noqa: E402 from pathocnn.model import build_model # noqa: E402 NUM_CLASSES = 9 def compute_split_stats(images_u8): """Per-channel mean/std with chunked float64 accumulation. A single float32 numpy reduction over ~90k images silently returns wrong statistics at this size (measured on this dataset: channel means of 0.23778/0.23778/0.23778 instead of 0.74055/0.53298/0.70583). Chunked float64 is the safe route; see REPORT.md. """ n = 0 s = np.zeros(3) ss = np.zeros(3) for i in range(0, len(images_u8), 10000): xb = images_u8[i:i + 10000].astype(np.float64) / 255.0 n += xb.shape[0] * 28 * 28 s += xb.sum(axis=(0, 1, 2)) ss += (xb ** 2).sum(axis=(0, 1, 2)) mean = s / n std = np.sqrt(ss / n - mean ** 2) return mean, std def normalize(x_u8_nchw): """uint8 NCHW tensor -> normalized float tensor (contract stats).""" x = x_u8_nchw.float() / 255.0 mean = torch.tensor(TRAIN_MEAN).view(1, 3, 1, 1) std = torch.tensor(TRAIN_STD).view(1, 3, 1, 1) return (x - mean) / std def augment(x_nchw, gen): """Random dihedral transform per sample (rot90 k in 0..3 + optional h-flip).""" n = x_nchw.shape[0] ks = torch.randint(0, 4, (n,), generator=gen) flips = torch.randint(0, 2, (n,), generator=gen) out = torch.empty_like(x_nchw) for k in range(4): for f in range(2): idx = (ks == k) & (flips == f) if idx.any(): t = torch.rot90(x_nchw[idx], k, dims=(2, 3)) if f: t = torch.flip(t, dims=(3,)) out[idx] = t return out @torch.no_grad() def evaluate(model, images_u8, labels, batch=512): """Return (probs np array, preds np array) for a split, contract preprocessing.""" model.eval() xt = torch.from_numpy(images_u8).permute(0, 3, 1, 2) probs = [] for i in range(0, len(xt), batch): xb = normalize(xt[i:i + batch]) probs.append(torch.softmax(model(xb), dim=1)) probs = torch.cat(probs).numpy() return probs, probs.argmax(1) def full_metrics(y_true, probs, preds): cm = confusion_matrix(y_true, preds, labels=list(range(NUM_CLASSES))) rec = recall_score(y_true, preds, labels=list(range(NUM_CLASSES)), average=None, zero_division=0) prec = precision_score(y_true, preds, labels=list(range(NUM_CLASSES)), average=None, zero_division=0) return { "accuracy": round(float(accuracy_score(y_true, preds)), 4), "auc_macro_ovr": round(float(roc_auc_score( y_true, probs, multi_class="ovr", average="macro")), 4), "balanced_accuracy": round(float(balanced_accuracy_score(y_true, preds)), 4), "confusion_matrix": cm.tolist(), "per_class_recall": {LABELS[i]: round(float(rec[i]), 4) for i in range(NUM_CLASSES)}, "per_class_precision": {LABELS[i]: round(float(prec[i]), 4) for i in range(NUM_CLASSES)}, "class_counts": np.bincount(y_true, minlength=NUM_CLASSES).tolist(), "majority_baseline": round(float(np.bincount(y_true).max() / len(y_true)), 4), "n": int(len(y_true)), } def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", required=True, help="path to pathmnist.npz") ap.add_argument("--out", required=True, help="work dir for checkpoint + metrics") ap.add_argument("--epochs", type=int, default=12) ap.add_argument("--batch", type=int, default=256) ap.add_argument("--lr", type=float, default=2e-3) ap.add_argument("--seed", type=int, default=42) args = ap.parse_args() os.makedirs(args.out, exist_ok=True) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) gen = torch.Generator().manual_seed(args.seed) d = np.load(args.data) x_train_u8 = d["train_images"] y_train = d["train_labels"].ravel().astype(np.int64) x_val_u8 = d["val_images"] y_val = d["val_labels"].ravel().astype(np.int64) x_test_u8 = d["test_images"] y_test = d["test_labels"].ravel().astype(np.int64) exact_mean, exact_std = compute_split_stats(x_train_u8) print(f"train stats (chunked float64): mean={exact_mean.round(5).tolist()} " f"std={exact_std.round(5).tolist()}", flush=True) print(f"contract stats used: mean={list(TRAIN_MEAN)} std={list(TRAIN_STD)}", flush=True) assert np.abs(exact_mean - np.array(TRAIN_MEAN)).max() < 1e-4 assert np.abs(exact_std - np.array(TRAIN_STD)).max() < 1e-4 xt = torch.from_numpy(x_train_u8).permute(0, 3, 1, 2) # uint8 NCHW in RAM yt = torch.from_numpy(y_train) model = build_model() n_params = sum(p.numel() for p in model.parameters()) print(f"model parameters: {n_params}", flush=True) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) steps_per_epoch = (len(xt) + args.batch - 1) // args.batch sched = torch.optim.lr_scheduler.OneCycleLR( opt, max_lr=args.lr, total_steps=steps_per_epoch * args.epochs, pct_start=0.1, anneal_strategy="cos") lossf = nn.CrossEntropyLoss() log = [] best_val_acc = -1.0 best_path = os.path.join(args.out, "pathocnn.pt") t_start = time.time() for epoch in range(args.epochs): model.train() perm = torch.randperm(len(xt), generator=gen) ep_loss, ep_correct, ep_n = 0.0, 0, 0 t0 = time.time() for i in range(0, len(xt), args.batch): idx = perm[i:i + args.batch] xb = augment(normalize(xt[idx]), gen) yb = yt[idx] opt.zero_grad() out = model(xb) loss = lossf(out, yb) loss.backward() opt.step() sched.step() ep_loss += loss.item() * len(idx) ep_correct += (out.argmax(1) == yb).sum().item() ep_n += len(idx) val_probs, val_preds = evaluate(model, x_val_u8, y_val) val_acc = float(accuracy_score(y_val, val_preds)) val_auc = float(roc_auc_score(y_val, val_probs, multi_class="ovr", average="macro")) row = { "epoch": epoch + 1, "train_loss": round(ep_loss / ep_n, 4), "train_acc": round(ep_correct / ep_n, 4), "val_acc": round(val_acc, 4), "val_auc_macro_ovr": round(val_auc, 4), "epoch_seconds": round(time.time() - t0, 1), } log.append(row) print(json.dumps(row), flush=True) if val_acc > best_val_acc: best_val_acc = val_acc torch.save({ "state_dict": model.state_dict(), "epoch": epoch + 1, "val_acc": val_acc, "val_auc_macro_ovr": val_auc, "seed": args.seed, "train_mean": list(TRAIN_MEAN), "train_std": list(TRAIN_STD), }, best_path) print(f" saved best checkpoint (val_acc {val_acc:.4f})", flush=True) total_seconds = round(time.time() - t_start, 1) # Final: load the best-val checkpoint and evaluate the test split once. ckpt = torch.load(best_path, map_location="cpu", weights_only=False) model.load_state_dict(ckpt["state_dict"]) test_probs, test_preds = evaluate(model, x_test_u8, y_test) metrics = { "model": "PathoCNN", "version": "0.1.0", "dataset": "PathMNIST (MedMNIST v2), official splits", "seed": args.seed, "epochs_run": args.epochs, "batch": args.batch, "lr": args.lr, "best_epoch": ckpt["epoch"], "best_val_acc": round(float(ckpt["val_acc"]), 4), "best_val_auc_macro_ovr": round(float(ckpt["val_auc_macro_ovr"]), 4), "train_seconds_total": total_seconds, "train_mean_exact": exact_mean.round(5).tolist(), "train_std_exact": exact_std.round(5).tolist(), "test": full_metrics(y_test, test_probs, test_preds), "epoch_log": log, } with open(os.path.join(args.out, "pathocnn_metrics.json"), "w") as f: json.dump(metrics, f, indent=2) print(json.dumps(metrics["test"], indent=2), flush=True) print(f"DONE total_seconds={total_seconds} best_epoch={ckpt['epoch']} " f"best_val_acc={best_val_acc:.4f}", flush=True) if __name__ == "__main__": raise SystemExit(main())