Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/train_asd.py

@@ -0,0 +1,180 @@
+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()