| """Recompute a run's val/test probabilities in dataset order, on one GPU. |
| |
| train.py evaluates under DDP, which returns rows in rank-interleaved order with |
| a padded tail. Combining runs needs rows aligned by study id, so this script |
| re-scores val and test on one GPU with a sequential loader and saves the uids. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| import sys |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from dataset import AtlasDataset, ShardIndex, to_model_input |
| from model import build_model |
| from train import build_splits, load_tables |
|
|
|
|
| @torch.no_grad() |
| def infer(model, loader, device, tta: bool = True): |
| P = [] |
| for xu, _y, pres in loader: |
| x = to_model_input(xu.to(device, non_blocking=True), |
| pres.to(device, non_blocking=True)) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| lg = model(x).float() |
| if tta: |
| lg = (lg + model(torch.flip(x, dims=[2])).float()) / 2.0 |
| P.append(torch.sigmoid(lg).cpu()) |
| return torch.cat(P).numpy().astype(np.float64) |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--run", required=True) |
| ap.add_argument("--cache", required=True) |
| ap.add_argument("--labels", required=True) |
| ap.add_argument("--splits", required=True) |
| ap.add_argument("--batch-size", type=int, default=48) |
| ap.add_argument("--workers", type=int, default=8) |
| ap.add_argument("--no-tta", action="store_true") |
| args = ap.parse_args() |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| ck = torch.load(os.path.join(args.run, "best.pt"), map_location="cpu", weights_only=False) |
| model = build_model(ck["arch"], num_classes=len(ck["labels"]), |
| in_ch=int(ck["model"]["stem.0.weight"].shape[1])) |
| model.load_state_dict(ck["model"]) |
| model = model.to(device).eval() |
|
|
| index = ShardIndex(args.cache) |
| labels, split, cols, _scored = load_tables(args.labels, args.splits) |
| if list(cols) != list(ck["labels"]): |
| raise SystemExit("label columns differ from the checkpoint's") |
| sp = build_splits(index, labels, split, len(cols)) |
|
|
| for name in ("val", "test"): |
| ids, y = sp[name] |
| ds = AtlasDataset(index, ids, y, ck.get("geometry", "d15"), train=False) |
| dl = DataLoader(ds, batch_size=args.batch_size, shuffle=False, |
| num_workers=args.workers, pin_memory=True) |
| probs = infer(model, dl, device, tta=not args.no_tta) |
| uids = np.array([index.uids[i] for i in ids]) |
| assert len(probs) == len(uids) == len(y) |
| out = os.path.join(args.run, f"{name}_probs_fixed.npz") |
| np.savez_compressed(out, probs=probs, y=y, uids=uids) |
| print(f"[{os.path.basename(args.run)}] {name}: {len(uids)} studies -> {out}", |
| flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
