"""Modality-specific foundation-model encoders for the H&E / mIF / markers / text binding (ImageBind-style). Same trainer as 09_bench.py, but every modality can use its own *pretrained, modality-specific* foundation model (frozen; features precomputed) followed by a trainable MLP head: --if_enc fm: -> cache/if_.npy (N,8,D; 8 dihedral views, aligned with the H&E views) e.g. fm:virtues --mk_enc fm: -> cache/markers_.npy (N,1,D) e.g. fm:geneformer, fm:scgpt --txt_enc -> cache/text_.npy (N,4,D) e.g. pathgenclip --save_heads -> runs/bench_fm//heads.pt (best-epoch encoders + input normalisation, for 31_virtual_validation.py) The from-scratch encoders of 09_bench.py (resnet18 | convnext_tiny | chanset, linear | mlp | ftt) remain available as baselines. Runs are written to runs/bench_fm/. Original 09 description: With a frozen anchor and anchor-only losses, each trainable encoder receives gradient only from its own InfoNCE(he, m) term, so one run trains one IF encoder, one marker encoder and one text head *independently*; each modality keeps its own best epoch (val R@1 of he<->m) and its test embeddings are saved at that epoch. Cross pairs (if<->markers, if<->text, markers<->text) are then evaluated post hoc for any encoder combination (10_bench_summary.py). IF encoders : resnet18 | convnext_tiny | chanset (shared per-channel ResNet18 + channel tokens + transformer) marker encoders : linear | mlp | ftt (FT-Transformer, one token per marker) text encoders : plip | biomedclip | imagebind (frozen text tower, trainable MLP head) """ import argparse import json import os import time import numpy as np import pandas as pd import timm import torch import torch.nn as nn import torch.nn.functional as F from common import CACHE, RUNS, SPLITS, dihedral p = argparse.ArgumentParser() p.add_argument("--split", default="orion_official") p.add_argument("--fold", type=int, default=0) p.add_argument("--anchor", default="uni2") p.add_argument("--if_enc", default="resnet18", help="resnet18|convnext_tiny|chanset|fm:") p.add_argument("--mk_enc", default="mlp", help="linear|mlp|ftt|fm:") p.add_argument("--txt_enc", default="plip", help="any with cache/text_.npy") p.add_argument("--epochs", type=int, default=40) p.add_argument("--bs", type=int, default=128) p.add_argument("--lr", type=float, default=3e-4) p.add_argument("--seed", type=int, default=0) p.add_argument("--save_heads", action="store_true", help="also save each modality's encoder at its best epoch (heads.pt)") args = p.parse_args() torch.manual_seed(args.seed); np.random.seed(args.seed) dev = "cuda" TRAIN = ["if", "markers", "text"] ENC = {"if": args.if_enc, "markers": args.mk_enc, "text": args.txt_enc} tag = lambda e: e.replace("fm:", "fm-") run_dir = os.path.join(RUNS, "bench_fm", f"{args.split}_fold{args.fold}_{args.anchor}_{tag(args.if_enc)}_{tag(args.mk_enc)}_{args.txt_enc}") os.makedirs(run_dir, exist_ok=True) # ---------------- data ---------------- sp = pd.read_csv(os.path.join(SPLITS, args.split, f"fold{args.fold}.csv")) HE = torch.from_numpy(np.load(os.path.join(CACHE, f"he_{args.anchor}.npy")).astype(np.float32)) # N,8,D D = HE.shape[-1] IF_FM, MK_FM = args.if_enc.startswith("fm:"), args.mk_enc.startswith("fm:") if IF_FM: # frozen mIF foundation-model features, one per dihedral view IF = torch.from_numpy(np.load(os.path.join(CACHE, f"if_{args.if_enc[3:]}.npy")).astype(np.float32)) else: IF = torch.from_numpy(np.load(os.path.join(CACHE, "if.npy"))) if MK_FM: # frozen gene-expression foundation-model features of the marker vector MK = torch.from_numpy(np.load(os.path.join(CACHE, f"markers_{args.mk_enc[3:]}.npy"))[:, 0].astype(np.float32)) else: MK = torch.from_numpy(np.load(os.path.join(CACHE, "markers.npy"))) TX = torch.from_numpy(np.load(os.path.join(CACHE, f"text_{args.txt_enc}.npy")).astype(np.float32)) # N,4,Dt tr_idx, va_idx, te_idx = (sp.tile_id[sp.split == s].values for s in ["train", "val", "test"]) if IF_FM: if_mean, if_std = IF[tr_idx].mean((0, 1)).to(dev), IF[tr_idx].std((0, 1)).to(dev) + 1e-3 else: if_mean = IF[tr_idx].float().mean((0, 2, 3)).view(1, -1, 1, 1).to(dev) / 255 if_std = IF[tr_idx].float().std((0, 2, 3)).view(1, -1, 1, 1).to(dev) / 255 + 1e-3 mk_mean, mk_std = MK[tr_idx].mean(0), MK[tr_idx].std(0) + 1e-3 def batch(idx, train=False): idx = torch.as_tensor(idx) d = torch.randint(0, 8, (len(idx),)) if train else torch.zeros(len(idx), dtype=torch.long) v = torch.randint(0, 4, (len(idx),)) if train else torch.zeros(len(idx), dtype=torch.long) if IF_FM: # precomputed features of the same dihedral view as the H&E anchor x = IF[idx, d].to(dev) else: x = IF[idx].to(dev, non_blocking=True).float() / 255 x = torch.stack([dihedral(a, int(k)) for a, k in zip(x, d)]) if train: x = x * (1 + 0.1 * torch.randn(len(idx), x.shape[1], 1, 1, device=dev)) mk = (MK[idx] - mk_mean) / mk_std if train and not MK_FM: mk = mk + 0.05 * torch.randn_like(mk) return {"he": HE[idx, d].to(dev), "if": (x - if_mean) / if_std, "markers": mk.to(dev), "text": TX[idx, v].to(dev)} # ---------------- encoders ---------------- def head(din): return nn.Sequential(nn.LayerNorm(din), nn.Linear(din, D)) class TimmEnc(nn.Module): def __init__(self, name): super().__init__() tag = {"resnet18": "resnet18.a1_in1k", "convnext_tiny": "convnext_tiny.fb_in1k"}[name] self.trunk = timm.create_model(tag, pretrained=True, in_chans=17, num_classes=0) self.head = head(self.trunk.num_features) def forward(self, x): return self.head(self.trunk(x)) class ChanSet(nn.Module): """Channel-agnostic mIF encoder: each channel -> shared 1-ch ResNet18 (112px) -> token (+ channel embedding) -> 2-layer transformer over channel tokens -> mean pool.""" def __init__(self, C=17, d=256): super().__init__() self.cnn = timm.create_model("resnet18.a1_in1k", pretrained=True, in_chans=1, num_classes=0) self.proj = nn.Linear(512, d) self.ch = nn.Parameter(torch.zeros(1, C, d)); nn.init.normal_(self.ch, std=0.02) self.tf = nn.TransformerEncoder(nn.TransformerEncoderLayer(d, 4, 4 * d, 0.1, batch_first=True, norm_first=True), 2) self.head = head(d) def forward(self, x): B, C, H, W = x.shape x = F.interpolate(x, size=(112, 112), mode="bilinear", antialias=True) t = self.proj(self.cnn(x.reshape(B * C, 1, 112, 112))).view(B, C, -1) + self.ch return self.head(self.tf(t).mean(1)) class FTT(nn.Module): """FT-Transformer: feature tokenizer (value * w_j + b_j) + CLS, 3 layers.""" def __init__(self, F_in, d=128): super().__init__() self.w = nn.Parameter(torch.randn(F_in, d) * 0.02) self.b = nn.Parameter(torch.zeros(F_in, d)) self.cls = nn.Parameter(torch.zeros(1, 1, d)) self.tf = nn.TransformerEncoder(nn.TransformerEncoderLayer(d, 8, 2 * d, 0.1, batch_first=True, norm_first=True), 3) self.head = head(d) def forward(self, x): t = x.unsqueeze(-1) * self.w + self.b t = torch.cat([self.cls.expand(len(x), -1, -1), t], 1) return self.head(self.tf(t)[:, 0]) def mlp(din, 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, D)) enc = nn.ModuleDict({ "if": mlp(IF.shape[-1]) if IF_FM else ChanSet() if args.if_enc == "chanset" else TimmEnc(args.if_enc), "markers": mlp(MK.shape[1]) if MK_FM else {"linear": lambda: nn.Linear(MK.shape[1], D), "mlp": lambda: mlp(MK.shape[1]), "ftt": lambda: FTT(MK.shape[1])}[args.mk_enc](), "text": mlp(TX.shape[-1]), # frozen text tower features -> trainable projection }).to(dev) logit = nn.ParameterDict({m: nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32)) for m in TRAIN}).to(dev) params = list(enc.parameters()) + list(logit.parameters()) opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.05) steps = args.epochs * (len(tr_idx) // args.bs) sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05) def embed(b): z = {"he": F.normalize(b["he"], dim=-1)} for m in TRAIN: z[m] = F.normalize(enc[m](b[m]).float(), dim=-1) return z @torch.no_grad() def embed_all(idx): enc.eval() out = {m: [] for m in ["he"] + TRAIN} for s in range(0, len(idx), 256): with torch.autocast("cuda", dtype=torch.bfloat16): z = embed(batch(idx[s:s + 256])) for m in out: out[m].append(z[m].float().cpu()) return {m: torch.cat(v) for m, v in out.items()} def r1(a, b): s = a @ b.T return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()) # ---------------- train with per-modality model selection ---------------- best = {m: -1.0 for m in TRAIN} best_ep = {m: -1 for m in TRAIN} test_z = {} best_state = {} hist = [] t0 = time.time() for ep in range(args.epochs): enc.train() perm = np.random.permutation(tr_idx) tl = {m: [] for m in TRAIN} for s in range(0, len(perm) - args.bs + 1, args.bs): b = batch(perm[s:s + args.bs], train=True) with torch.autocast("cuda", dtype=torch.bfloat16): z = embed(b) tgt = torch.arange(args.bs, device=dev) loss = 0 for m in TRAIN: lg = logit[m].exp().clamp(max=100) * z["he"] @ z[m].T l = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2 loss = loss + l tl[m].append(l.item()) opt.zero_grad(set_to_none=True) loss.backward() opt.step(); sched.step() zv = embed_all(va_idx) rec = {"epoch": ep} improved = [] for m in TRAIN: v = (r1(zv["he"], zv[m]) + r1(zv[m], zv["he"])) / 2 rec[f"loss_{m}"], rec[f"val_r1_{m}"] = float(np.mean(tl[m])), v if v > best[m]: best[m], best_ep[m] = v, ep improved.append(m) if improved: zt = embed_all(te_idx) test_z["he"] = zt["he"] for m in improved: test_z[m] = zt[m] if args.save_heads: best_state[m] = {k: v.detach().cpu().clone() for k, v in enc[m].state_dict().items()} hist.append(rec) print(f"ep {ep} " + " ".join(f"{m}:{rec[f'val_r1_{m}']:.3f}" for m in TRAIN) + f" ({time.time() - t0:.0f}s)", flush=True) torch.save({"z": test_z, "tile_id": te_idx}, os.path.join(run_dir, "test_embeddings.pt")) if args.save_heads: # + the input normalisation needed to apply the heads to new data torch.save({"state": best_state, "if_mean": if_mean.cpu(), "if_std": if_std.cpu(), "mk_mean": mk_mean, "mk_std": mk_std}, os.path.join(run_dir, "heads.pt")) json.dump({"args": vars(args), "best_val": best, "best_epoch": best_ep, "n_train": len(tr_idx), "n_test": len(te_idx), "time_s": time.time() - t0, "history": hist}, open(os.path.join(run_dir, "train.json"), "w"), indent=1) print("done", best, best_ep)