icassp2027-experiment/third_party/MuVAP/train_asd.py

import argparse
from pathlib import Path

import torch
from lightning.pytorch import Trainer, seed_everything
from lightning.pytorch.callbacks import (
    Callback,
    LearningRateMonitor,
    ModelCheckpoint,
    RichProgressBar,
)

from config.setup import load_config
from data.packed_asd_dm import PackedASDDataModule
from models.release import load_weights, merge_config, resolve
from projection_window import ProjectionWindow
from tasks.asd_task import ASDTask


class BatchSamplerEpoch(Callback):
    """Advance the train batch sampler's epoch so its shuffle changes.

    Lightning only forwards `set_epoch` to `dataloader.sampler` and
    `dataloader.batch_sampler.sampler`, so a bare batch sampler never hears
    about the epoch and would replay one fixed batch order forever.
    """

    def on_train_epoch_start(self, trainer, pl_module):
        sampler = getattr(trainer.train_dataloader, "batch_sampler", None)
        set_epoch = getattr(sampler, "set_epoch", None)
        if callable(set_epoch):
            set_epoch(trainer.current_epoch)


def parse_args():
    parser = argparse.ArgumentParser(description="Train or evaluate ASD")
    parser.add_argument(
        "--config", default=Path(__file__).resolve().parent / "config/yaml/asd.yaml"
    )
    parser.add_argument("--test", action="store_true")
    parser.add_argument("--checkpoint", help="Lightning .ckpt to resume or evaluate")
    parser.add_argument(
        "--weights",
        help="published release directory (or its weights.pt) to evaluate; "
        "the model architecture comes from the release, data paths from --config",
    )
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--wandb", action="store_true")
    parser.add_argument(
        "--wandb-project",
        default="MuVAP-ASD",
        help="W&B project the run is logged to (with --wandb)",
    )
    parser.add_argument("--name", help="run name used for the checkpoint directory")
    parser.add_argument(
        "--set",
        action="append",
        default=[],
        metavar="KEY=VALUE",
        help="override one config value, e.g. --set asd.audio_encoder=mimi",
    )
    return parser.parse_args()


def apply_overrides(cfg, overrides):
    """Apply `--set a.b.c=value` onto a loaded config, parsed as YAML scalars."""
    import yaml

    for override in overrides:
        key, _, raw = override.partition("=")
        if not _:
            raise SystemExit(f"--set expects KEY=VALUE, got {override!r}")
        target = cfg
        *path, leaf = key.split(".")
        for step in path:
            if step not in target:
                raise SystemExit(f"unknown config key in --set {override!r}")
            target = target[step]
        if leaf not in target:
            raise SystemExit(f"unknown config key in --set {override!r}")
        target[leaf] = yaml.safe_load(raw)
    return cfg


def build_checkpoints(asd_cfg, checkpoint_dir):
    """Always keep `last.ckpt`; keep a selected `best.ckpt` only if asked to.

    Two callbacks, not one: a single `ModelCheckpoint` given both a monitor and
    `save_last=True` writes a `last.ckpt` that tracks the selected checkpoint
    rather than the final epoch, so the final epoch would never be kept.

    `val/mAP_official` is the reported metric; `val/loss` mixes the projection
    and contrastive terms and does not track it closely. Setting
    `checkpoint_monitor` to nothing says there is no held-out split to select
    on, and the final epoch becomes the model the run produced.
    """
    last = ModelCheckpoint(dirpath=checkpoint_dir, save_top_k=0, save_last=True)
    monitor = asd_cfg.get("checkpoint_monitor", "val/mAP_official")
    if not monitor:
        return [last]
    best = ModelCheckpoint(
        monitor=monitor,
        mode=asd_cfg.get("checkpoint_mode", "max"),
        save_top_k=1,
        dirpath=checkpoint_dir,
        filename="best",
    )
    return [best, last]


def main():
    args = parse_args()
    if args.test and not (args.checkpoint or args.weights):
        raise SystemExit("--test needs --checkpoint (a .ckpt) or --weights (a release)")

    cfg = apply_overrides(load_config(args.config, "asd"), args.set)
    weights_path = None
    if args.weights:
        weights_path, release_config = resolve(args.weights)
        cfg = merge_config(cfg, release_config, "asd")

    seed_everything(args.seed, workers=True)
    asd_cfg = cfg["asd"]
    # Ampere and newer trade a little float32 mantissa for a large matmul speedup.
    torch.set_float32_matmul_precision(asd_cfg.get("matmul_precision", "high"))

    logger = None
    run_id = args.name or "local"
    if args.wandb:
        from lightning.pytorch.loggers import WandbLogger

        logger = WandbLogger(project=args.wandb_project, name=args.name)
        run_id = args.name or logger.experiment.id

    checkpoint_dir = Path(cfg["checkpoint_dir"]) / run_id
    callbacks = [
        *build_checkpoints(asd_cfg, checkpoint_dir),
        BatchSamplerEpoch(),
        RichProgressBar(),
        LearningRateMonitor("step"),
    ]
    projection = ProjectionWindow(**cfg["asd"]["projection_window"])
    encoder = asd_cfg.get("audio_encoder", "cpc")
    print(f"ASD setup: audio_encoder={encoder} {projection}")

    data = PackedASDDataModule(cfg, projection)
    model = ASDTask(cfg)
    if weights_path is not None:
        load_weights(model.model, weights_path)
        print(f"loaded release weights from {weights_path}")
    trainer = Trainer(
        # Checkpoints and Lightning's fallback CSV logs both land here rather
        # than as `lightning_logs/` inside the code checkout.
        default_root_dir=checkpoint_dir,
        max_epochs=asd_cfg["max_epochs"],
        logger=logger,
        callbacks=callbacks,
        log_every_n_steps=100,
        accelerator="auto",
        devices="auto",
        precision=asd_cfg.get("precision", "bf16-mixed"),
        gradient_clip_val=asd_cfg.get("gradient_clip_val", 1.0),
        accumulate_grad_batches=asd_cfg.get("accumulate_grad_batches", 1),
        # The train loader owns a frame-budget batch sampler that already
        # shards across ranks; Lightning must not wrap it in its own sampler.
        use_distributed_sampler=False,
    )
    if args.test:
        # A release is already in the model; a ckpt_path would overwrite it.
        trainer.validate(
            model,
            datamodule=data,
            ckpt_path=None if weights_path is not None else args.checkpoint,
        )
    else:
        trainer.fit(model, datamodule=data)


if __name__ == "__main__":
    main()