scripts/train_pet_text_alignment.py

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()