import os import glob from typing import Any import click import lightning.pytorch as pl import torch from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint from lightning.pytorch.loggers import TensorBoardLogger from pytorchvideo.transforms import Normalize, Permute, RandAugment from torch.utils.data import DataLoader, WeightedRandomSampler from torchvision.transforms import transforms as T from torchvision.transforms._transforms_video import ToTensorVideo from torchvision.transforms import InterpolationMode from full_model.rnn_dataset import SyntaxDataset from full_model.rnn_model import SyntaxLightningModule torch.set_float32_matmul_precision("medium") def get_transforms(video_size, imagenet_mean, imagenet_std, train: bool = True): """Augmentations and preprocessing for clips.""" interpolation_choices = [InterpolationMode.BILINEAR, InterpolationMode.BICUBIC] if train: return T.Compose([ ToTensorVideo(), Permute(dims=[1, 0, 2, 3]), RandAugment(magnitude=10, num_layers=2), T.RandomHorizontalFlip(), Permute(dims=[1, 0, 2, 3]), T.RandomChoice([ T.Resize(size=video_size, interpolation=interp, antialias=True) for interp in interpolation_choices ]), Normalize(mean=imagenet_mean, std=imagenet_std), ]) else: return T.Compose([ ToTensorVideo(), T.Resize(size=video_size, interpolation=InterpolationMode.BICUBIC, antialias=True), Normalize(mean=imagenet_mean, std=imagenet_std), ]) def make_dataloader(dataset, batch_size: int, num_workers: int, use_weighted_sampler: bool): """DataLoader with an optional score-based WeightedRandomSampler.""" if use_weighted_sampler: sample_weights = dataset.get_sample_weights().cpu() sampler = WeightedRandomSampler(sample_weights, num_samples=len(dataset), replacement=True) shuffle = False else: sampler = None shuffle = False return DataLoader( dataset, batch_size=batch_size, num_workers=num_workers, sampler=sampler, shuffle=shuffle, drop_last=True, pin_memory=True, persistent_workers=(num_workers > 0), ) def make_model( num_classes: int, lr: float, variant: str, weight_decay: float, max_epochs: int, weight_path: str | None = None, pl_weight_path: str | None = None, pt_weights_format: bool = False, ) -> SyntaxLightningModule: """ Create the head model. weight_path — pretrained backbone weights (r3d_18), .pt or .ckpt. pl_weight_path — full head-model checkpoint (Lightning .ckpt or raw .pt). pt_weights_format=True → pl_weight_path is a raw state_dict (.pt). pt_weights_format=False → pl_weight_path is a Lightning .ckpt with 'state_dict'. """ return SyntaxLightningModule( num_classes=num_classes, lr=lr, variant=variant, weight_decay=weight_decay, max_epochs=max_epochs, weight_path=weight_path, pl_weight_path=pl_weight_path, yulie_model=pt_weights_format, ) def make_callbacks(phase: str): """Callbacks: LR monitor plus ModelCheckpoint on val_rmse.""" lr_monitor = LearningRateMonitor(logging_interval="epoch") if phase == "pre": checkpoint = ModelCheckpoint( monitor="val_rmse", save_top_k=1, mode="min", filename="rnn_model-{epoch:02d}-{val_rmse:.3f}", save_last=True, ) elif phase == "full": checkpoint = ModelCheckpoint( monitor="val_rmse", save_top_k=3, mode="min", filename="rnn_model-{epoch:02d}-{val_rmse:.3f}", save_last=True, ) else: raise ValueError(f"Unknown phase '{phase}', expected 'pre' or 'full'") return [lr_monitor, checkpoint] def make_trainer(max_epochs: int, logdir: str, logger_name: str, devices: list[int], precision: str, callbacks): """Create a Trainer with a TensorBoard logger.""" logger = TensorBoardLogger(save_dir=logdir, name=logger_name) strategy = "ddp_find_unused_parameters_true" if len(devices) > 1 else "auto" trainer = pl.Trainer( max_epochs=max_epochs, accelerator="gpu" if torch.cuda.is_available() else "cpu", devices=devices, strategy=strategy, precision=precision, callbacks=callbacks, log_every_n_steps=10, logger=logger, ) return trainer def find_backbone_ckpt_lightning(backbone_logdir: str, artery: str, fold: int, phase: str = "full") -> str: """ Find a backbone Lightning checkpoint in the log directory. Expected structure: backbone_logdir/ {artery}BinSyntax_R3D_{phase}_foldXX/version_*/checkpoints/*.ckpt """ logger_name = f"{artery}BinSyntax_R3D_{phase}_fold{fold:02d}" pattern = os.path.join(backbone_logdir, logger_name, "version_*/checkpoints", "*.ckpt") ckpts = glob.glob(pattern) if not ckpts: raise FileNotFoundError( f"No backbone Lightning checkpoints found for\n" f" artery={artery}, fold={fold}, phase={phase}\n" f" in '{backbone_logdir}' (pattern: {pattern})" ) best = max(ckpts, key=os.path.getctime) print(f"[Backbone] Using Lightning checkpoint: {best}") return best def build_backbone_pt_path(backbone_pt_dir: str, artery: str, fold: int) -> str: """ Build the backbone .pt path using the naming convention: rightBinSyntax_R3D_full_fold00.pt leftBinSyntax_R3D_full_fold00.pt ... """ fname = f"{artery}BinSyntax_R3D_full_fold{fold:02d}.pt" path = os.path.join(backbone_pt_dir, fname) if not os.path.exists(path): raise FileNotFoundError( f"Backbone .pt not found for artery={artery}, fold={fold} in '{backbone_pt_dir}'\n" f"Expected file: {fname}" ) print(f"[Backbone] Using .pt file: {path}") return path @click.command() @click.option( "-r", "--dataset-root", type=click.Path(exists=True), default=".", show_default=True, help="Dataset root (JSON and DICOM paths are resolved relative to it).", ) @click.option("--fold", type=int, default=4, show_default=True, help="Fold number.") @click.option( "-a", "--artery", type=str, default="right", show_default=True, help="Artery: left or right.", ) @click.option( "--variant", type=str, default="lstm_mean", show_default=True, help="Head-model variant: mean_out, mean, lstm_mean, lstm_last, gru_mean, gru_last, bert_mean, bert_cls, bert_cls2.", ) @click.option("-nc", "--num-classes", type=int, default=2, show_default=True, help="Number of head-model outputs (clf + reg).") @click.option("-b", "--batch-size", type=int, default=8, show_default=True, help="Batch size.") @click.option("-f", "--frames-per-clip", type=int, default=32, show_default=True, help="Frames per clip.") @click.option( "-v", "--video-size", type=click.Tuple([int, int]), default=(256, 256), show_default=True, help="Frame size (H, W).", ) @click.option("--max-epochs", type=int, default=10, show_default=True, help="Number of full-train epochs.") @click.option("--num-workers", type=int, default=16, show_default=True, help="DataLoader workers.") @click.option( "--devices", type=list[int], multiple=True, default=[0], show_default=True, help="List of GPU ids", ) @click.option("--precision", type=str, default="bf16-mixed", show_default=True, help="Numeric precision mode.") @click.option( "--logdir", type=click.Path(), default="./logs/rnn", show_default=True, help="Log and checkpoint directory for the head model.", ) @click.option( "--backbone-logdir", type=click.Path(exists=True), default=None, help="Directory with backbone logs (Lightning .ckpt files).", ) @click.option( "--backbone-pt-dir", type=click.Path(exists=True), default="backbone_weights", show_default=True, help="Directory with backbone .pt files (rightBinSyntax_R3D_full_foldXX.pt, leftBinSyntax_R3D_full_foldXX.pt).", ) @click.option( "--backbone-from-pt", is_flag=True, default=True, show_default=True, help="When enabled, load the backbone from .pt files in backbone-pt-dir; otherwise use Lightning logs in backbone-logdir.", ) @click.option( "--rnn-folds-dir", type=click.Path(), default="rnn_folds", show_default=True, help="Directory with rnn_folds (relative to dataset_root).", ) @click.option( "--use-weighted-sampler", is_flag=True, default=False, show_default=True, help="Use a WeightedRandomSampler by score.", ) @click.option( "--pt-weights-format", is_flag=True, default=False, show_default=True, help="pl_weight_path format for full training: True uses .pt raw state_dict, False uses Lightning .ckpt.", ) @click.option("--seed", type=int, default=42, show_default=True, help="Random seed.") def main( dataset_root: str, fold: int, artery: str, variant: str, num_classes: int, batch_size: int, frames_per_clip: int, video_size: Any, max_epochs: int, num_workers: int, devices: int, precision: str, logdir: str, backbone_logdir: str | None, backbone_pt_dir: str | None, backbone_from_pt: bool, rnn_folds_dir: str, use_weighted_sampler: bool, pt_weights_format: bool, seed: int, ): """Train the RNN head on top of the backbone.""" VARIANTS = "mean_out mean lstm_mean lstm_last gru_mean gru_last bert_mean bert_cls bert_cls2".split() if variant not in VARIANTS: raise ValueError(f"Unknown variant '{variant}', expected one of: {VARIANTS}") artery = artery.lower() if artery not in ("left", "right"): raise ValueError(f"Unknown artery '{artery}', expected 'left' or 'right'") pl.seed_everything(seed) imagenet_mean = [0.485, 0.456, 0.406] imagenet_std = [0.229, 0.224, 0.225] train_meta = os.path.join(rnn_folds_dir, f"rnn_fold{fold:02d}_train.json") eval_meta = os.path.join(rnn_folds_dir, f"rnn_fold{fold:02d}_eval.json") train_set = SyntaxDataset( root=dataset_root, meta=train_meta, train=True, length=frames_per_clip, label=f"syntax_{artery}", artery=artery, inference=False, validation=True, transform=get_transforms(video_size, imagenet_mean, imagenet_std, train=True), ) val_set = SyntaxDataset( root=dataset_root, meta=eval_meta, train=False, length=frames_per_clip, label=f"syntax_{artery}", artery=artery, inference=False, validation=True, transform=get_transforms(video_size, imagenet_mean, imagenet_std, train=False), ) train_loader_pre = make_dataloader(train_set, batch_size * 2, num_workers, use_weighted_sampler) train_loader_post = make_dataloader(train_set, batch_size, num_workers, use_weighted_sampler) val_loader = make_dataloader(val_set, 1, num_workers, use_weighted_sampler=False) x, *_ = next(iter(train_loader_pre)) video_shape = x.shape[2:] print(f"RNN head input per clip: {video_shape}") if backbone_from_pt: if backbone_pt_dir is None: raise ValueError("backbone-from-pt=True, but backbone-pt-dir is not set.") backbone_weight_path = build_backbone_pt_path(backbone_pt_dir, artery=artery, fold=fold) else: if backbone_logdir is None: raise ValueError("backbone-from-pt=False, but backbone-logdir is not set.") backbone_weight_path = find_backbone_ckpt_lightning( backbone_logdir=backbone_logdir, artery=artery, fold=fold, phase="full", ) callbacks_pre = make_callbacks(phase="pre") model_pre = make_model( num_classes=num_classes, lr=1e-4, variant=variant, weight_decay=0.01, max_epochs=max_epochs, weight_path=backbone_weight_path, pl_weight_path=None, pt_weights_format=False, ) trainer_pre = make_trainer( max_epochs=max_epochs, logdir=logdir, logger_name=f"{artery}BinSyntax_R3D_fold{fold:02d}_{variant}_pre", devices=devices, precision=precision, callbacks=callbacks_pre, ) trainer_pre.fit(model_pre, train_dataloaders=train_loader_pre, val_dataloaders=val_loader) callbacks_full = make_callbacks(phase="full") model_full = make_model( num_classes=num_classes, lr=2e-5, variant=variant, weight_decay=0.01, max_epochs=max_epochs, weight_path=None, pl_weight_path=trainer_pre.checkpoint_callback.best_model_path, pt_weights_format=pt_weights_format, ) trainer_full = make_trainer( max_epochs=max_epochs, logdir=logdir, logger_name=f"{artery}BinSyntax_R3D_fold{fold:02d}_{variant}_post", devices=devices, precision=precision, callbacks=callbacks_full, ) trainer_full.fit(model_full, train_dataloaders=train_loader_post, val_dataloaders=val_loader) if __name__ == "__main__": main()