| """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:<name> -> cache/if_<name>.npy (N,8,D; 8 dihedral views, aligned with the H&E views) e.g. fm:virtues |
| --mk_enc fm:<name> -> cache/markers_<name>.npy (N,1,D) e.g. fm:geneformer, fm:scgpt |
| --txt_enc <name> -> cache/text_<name>.npy (N,4,D) e.g. pathgenclip |
| --save_heads -> runs/bench_fm/<run>/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:<name>") |
| p.add_argument("--mk_enc", default="mlp", help="linear|mlp|ftt|fm:<name>") |
| p.add_argument("--txt_enc", default="plip", help="any <name> with cache/text_<name>.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) |
|
|
| |
| 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)) |
| D = HE.shape[-1] |
| IF_FM, MK_FM = args.if_enc.startswith("fm:"), args.mk_enc.startswith("fm:") |
| if IF_FM: |
| 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: |
| 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)) |
| 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: |
| 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)} |
|
|
|
|
| |
| 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]), |
| }).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()) |
|
|
|
|
| |
| 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: |
| 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) |
|
|
