| """Bind real pathology text to the frozen H&E anchor using PathGen-1.6M H&E-caption pairs (ImageBind recipe: |
| text is paired only with images; it never sees ORION, IF or markers). |
| |
| frozen text tower features (cache/pathgen_text_<enc>.npy) -> trainable MLP head -> anchor space |
| loss: symmetric InfoNCE vs frozen anchor image features (cache/pathgen_he_<anchor>.npy, 8 dihedral views) |
| Usage: python 12_train_text_pathgen.py --txt_enc plip|biomedclip|imagebind [--anchor uni2] |
| Saves runs/text_pathgen/<anchor>_<enc>/{head.pt, metrics.json} |
| """ |
| import argparse |
| import json |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from common import CACHE, RUNS |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--txt_enc", default="plip") |
| p.add_argument("--anchor", default="uni2") |
| p.add_argument("--epochs", type=int, default=100) |
| p.add_argument("--bs", type=int, default=256) |
| p.add_argument("--lr", type=float, default=3e-4) |
| p.add_argument("--seed", type=int, default=0) |
| args = p.parse_args() |
| torch.manual_seed(args.seed); np.random.seed(args.seed) |
| dev = "cuda" |
| out_dir = os.path.join(RUNS, "text_pathgen", f"{args.anchor}_{args.txt_enc}") |
| os.makedirs(out_dir, exist_ok=True) |
|
|
| meta = pd.read_csv(os.path.join(CACHE, "pathgen_meta.csv")) |
| HE = torch.from_numpy(np.load(os.path.join(CACHE, f"pathgen_he_{args.anchor}.npy")).astype(np.float32)) |
| TX = torch.from_numpy(np.load(os.path.join(CACHE, f"pathgen_text_{args.txt_enc}.npy"))[:, 0].astype(np.float32)) |
| D = HE.shape[-1] |
| idx = {s: meta.idx[meta.split == s].values for s in ["train", "val", "test"]} |
|
|
|
|
| def make_head(din, dout, hid=512): |
| return nn.Sequential(nn.Linear(din, hid), nn.GELU(), nn.Dropout(0.1), nn.Linear(hid, hid), nn.GELU(), |
| nn.LayerNorm(hid), nn.Linear(hid, dout)) |
|
|
|
|
| head = make_head(TX.shape[1], D).to(dev) |
| logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev)) |
| opt = torch.optim.AdamW(list(head.parameters()) + [logit], lr=args.lr, weight_decay=0.05) |
| steps = args.epochs * max(1, len(idx["train"]) // args.bs) |
| sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05) |
|
|
|
|
| @torch.no_grad() |
| def retrieval(ids, size=None, reps=50): |
| head.eval() |
| zi = F.normalize(HE[ids, 0], dim=-1) |
| zt = F.normalize(head(TX[ids].to(dev)).float().cpu(), dim=-1) |
|
|
| def r1(a, b): |
| s = a @ b.T |
| return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()) |
|
|
| if size is None or len(ids) <= size: |
| return {"he->text": r1(zi, zt), "text->he": r1(zt, zi), "N": len(ids)} |
| rng = np.random.default_rng(0) |
| a, b = [], [] |
| for _ in range(reps): |
| s = torch.from_numpy(rng.choice(len(ids), size, replace=False)) |
| a.append(r1(zi[s], zt[s])); b.append(r1(zt[s], zi[s])) |
| return {"he->text": float(np.mean(a)), "text->he": float(np.mean(b)), "N": size} |
|
|
|
|
| best, hist = -1, [] |
| for ep in range(args.epochs): |
| head.train() |
| perm = np.random.permutation(idx["train"]) |
| for s in range(0, len(perm) - args.bs + 1, args.bs): |
| b = perm[s:s + args.bs] |
| zi = F.normalize(HE[b, torch.randint(0, 8, (len(b),))].to(dev), dim=-1) |
| zt = F.normalize(head(TX[b].to(dev)), dim=-1) |
| lg = logit.exp().clamp(max=100) * zi @ zt.T |
| tgt = torch.arange(len(b), device=dev) |
| loss = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2 |
| opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step() |
| v = retrieval(idx["val"]) |
| vm = (v["he->text"] + v["text->he"]) / 2 |
| hist.append({"epoch": ep, "loss": loss.item(), "val_r1": vm}) |
| if vm > best: |
| best = vm |
| torch.save(head.state_dict(), os.path.join(out_dir, "head.pt")) |
|
|
| head.load_state_dict(torch.load(os.path.join(out_dir, "head.pt"))) |
| res = {"args": vars(args), "best_val_r1": best, "test_full": retrieval(idx["test"]), |
| "test_g100": retrieval(idx["test"], size=100), "n": {k: len(v) for k, v in idx.items()}, "history": hist} |
| |
| zs_img = os.path.join(CACHE, f"pathgen_he_{args.txt_enc}.npy") |
| if os.path.exists(zs_img): |
| zi = F.normalize(torch.from_numpy(np.load(zs_img)[idx["test"], 0].astype(np.float32)), dim=-1) |
| zt = F.normalize(TX[idx["test"]], dim=-1) |
| s = zi @ zt.T |
| res["zeroshot_test_full"] = {"he->text": float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()), |
| "text->he": float(((s.T > s.diag()[:, None]).sum(1) == 0).float().mean())} |
| json.dump(res, open(os.path.join(out_dir, "metrics.json"), "w"), indent=1) |
| print(args.txt_enc, "val", round(best, 3), "test full", res["test_full"], "g100", res["test_g100"]) |
|
|
