from __future__ import annotations import argparse from pathlib import Path import pandas as pd import torch from torch import nn from torch.utils.data import DataLoader, Dataset from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr from train_pet_foundation import PETSUVRFoundationModel, build_encoder class PETTextDataset(PETSUVRDataset): def __init__(self, text_csv: str | Path, output_size: tuple[int, int, int]) -> None: super().__init__(text_csv, output_size=output_size) def __getitem__(self, index: int) -> dict[str, object]: item = super().__getitem__(index) row = self.manifest.iloc[index] item["text"] = str(row["region_text"]) item["low_regions"] = str(row["low_regions"]) item["high_regions"] = str(row["high_regions"]) return item def collate_pet_text(batch: list[dict[str, object]]) -> dict[str, object]: out = collate_pet_suvr(batch) out["low_regions"] = [item["low_regions"] for item in batch] out["high_regions"] = [item["high_regions"] for item in batch] return out class PETTextAlignmentModel(nn.Module): def __init__(self, pet_model: PETSUVRFoundationModel, text_model_name: str, embed_dim: int) -> None: super().__init__() try: from transformers import AutoModel except ImportError as exc: raise ImportError("PET-text alignment requires `transformers`. Install it before training Stage 2.") from exc self.pet_model = pet_model for p in self.pet_model.parameters(): p.requires_grad = False self.pet_model.eval() self.text_model = AutoModel.from_pretrained(text_model_name) for p in self.text_model.parameters(): p.requires_grad = False self.text_model.eval() text_dim = int(getattr(self.text_model.config, "hidden_size")) self.pet_head = nn.Sequential(nn.LayerNorm(embed_dim), nn.Linear(embed_dim, embed_dim)) self.text_head = nn.Sequential(nn.LayerNorm(text_dim), nn.Linear(text_dim, embed_dim)) self.temperature = nn.Parameter(torch.tensor(0.07)) def encode_pet(self, image: torch.Tensor) -> torch.Tensor: with torch.no_grad(): pet_feat = self.pet_model.pet_encoder(image) pet_z = torch.nn.functional.normalize(self.pet_model.pet_projector(pet_feat), dim=-1) return torch.nn.functional.normalize(self.pet_head(pet_z), dim=-1) def encode_text(self, tokens: dict[str, torch.Tensor]) -> torch.Tensor: with torch.no_grad(): output = self.text_model(**tokens) pooled = getattr(output, "pooler_output", None) if pooled is None: mask = tokens["attention_mask"].unsqueeze(-1).to(output.last_hidden_state.dtype) pooled = (output.last_hidden_state * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) return torch.nn.functional.normalize(self.text_head(pooled), dim=-1) def forward(self, image: torch.Tensor, tokens: dict[str, torch.Tensor]) -> torch.Tensor: pet_z = self.encode_pet(image) text_z = self.encode_text(tokens) return pet_z @ text_z.T / self.temperature.clamp_min(0.01) def load_pet_model(checkpoint: Path, args: argparse.Namespace, device: torch.device) -> PETSUVRFoundationModel: ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) saved_args = ckpt.get("args", {}) for name in ("backbone", "embed_dim", "freeze_encoder"): if getattr(args, name, None) is None and name in saved_args: setattr(args, name, saved_args[name]) if args.output_size is None: args.output_size = tuple(saved_args.get("output_size", (96, 96, 96))) sample = PETSUVRDataset(args.train_csv, output_size=tuple(args.output_size))[0] encoder = build_encoder(args) model = PETSUVRFoundationModel(encoder, int(sample["suvr"].numel()), args.embed_dim or 256, bool(args.freeze_encoder)) model.load_state_dict(ckpt["model"], strict=True) return model.to(device) def run_epoch( model: PETTextAlignmentModel, loader: DataLoader, tokenizer, device: torch.device, optimizer: torch.optim.Optimizer | None, max_length: int, ) -> float: train = optimizer is not None model.train(train) model.pet_model.eval() model.text_model.eval() total = 0.0 count = 0 for batch in loader: image = batch["image"].to(device, non_blocking=True) tokens = tokenizer(batch["text"], padding=True, truncation=True, max_length=max_length, return_tensors="pt") tokens = {k: v.to(device) for k, v in tokens.items()} logits = model(image, tokens) labels = torch.arange(logits.shape[0], device=device) loss = 0.5 * ( nn.functional.cross_entropy(logits, labels) + nn.functional.cross_entropy(logits.T, labels) ) if train: optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() total += float(loss.detach()) * image.shape[0] count += image.shape[0] return total / max(count, 1) def main() -> None: parser = argparse.ArgumentParser(description="Train controlled PET-to-region-text contrastive alignment.") parser.add_argument("--pet-checkpoint", type=Path, required=True) parser.add_argument("--train-csv", type=Path, required=True) parser.add_argument("--val-csv", type=Path, required=True) parser.add_argument("--text-model", default="emilyalsentzer/Bio_ClinicalBERT") parser.add_argument("--backbone", choices=["small_cnn", "medicalnet", "brainiac", "brainfm", "swinunetr", "sam_med3d"], default=None) parser.add_argument("--medicalnet-weights", type=Path, default=Path("pretrained/medicalnet/resnet_50_23dataset.pth")) parser.add_argument("--brainiac-weights", type=Path, default=Path("pretrained/brainiac/backbone.safetensors")) parser.add_argument("--brainfm-weights", type=Path, default=Path("pretrained/brainfm/assets/brainfm_pretrained.pth")) parser.add_argument("--brainfm-code-root", type=Path, default=Path("pretrained/brainfm")) parser.add_argument("--swinunetr-weights", type=Path, default=Path("pretrained/swinunetr/model_swinvit.pt")) parser.add_argument("--sam-med3d-weights", type=Path, default=Path("pretrained/sam-med3d/sam_med3d_turbo.pth")) parser.add_argument("--output-size", type=int, nargs=3, default=None) parser.add_argument("--embed-dim", type=int, default=None) parser.add_argument("--freeze-encoder", action=argparse.BooleanOptionalAction, default=None) parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--num-workers", type=int, default=2) parser.add_argument("--max-length", type=int, default=96) parser.add_argument("--out", type=Path, default=Path("runs/vlm/pet_text_alignment.pt")) parser.add_argument("--best-out", type=Path, default=None) args = parser.parse_args() from transformers import AutoTokenizer device = torch.device("cuda" if torch.cuda.is_available() else "cpu") pet_model = load_pet_model(args.pet_checkpoint, args, device) tokenizer = AutoTokenizer.from_pretrained(args.text_model) model = PETTextAlignmentModel(pet_model, args.text_model, args.embed_dim or 256).to(device) train_set = PETTextDataset(args.train_csv, output_size=tuple(args.output_size)) val_set = PETTextDataset(args.val_csv, output_size=tuple(args.output_size)) train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=collate_pet_text) val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=collate_pet_text) optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr) best = float("inf") best_out = args.best_out or args.out.with_name(args.out.stem + "_best.pt") args.out.parent.mkdir(parents=True, exist_ok=True) for epoch in range(1, args.epochs + 1): train_loss = run_epoch(model, train_loader, tokenizer, device, optimizer, args.max_length) val_loss = run_epoch(model, val_loader, tokenizer, device, None, args.max_length) print(f"epoch={epoch} train_loss={train_loss:.6f} val_loss={val_loss:.6f}", flush=True) if val_loss < best: best = val_loss torch.save({"model": model.state_dict(), "args": vars(args), "best_val_loss": best, "epoch": epoch}, best_out) print(f"saved_best {best_out} val_loss={best:.6f}", flush=True) torch.save({"model": model.state_dict(), "args": vars(args), "best_val_loss": best}, args.out) print(f"saved {args.out}", flush=True) if __name__ == "__main__": main()