full_model/rnn_train.py

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