diff --git a/configs/callbacks/curtain_auc_percentiles.yaml b/configs/callbacks/curtain_auc_percentiles.yaml new file mode 100644 index 0000000..8c9c0f4 --- /dev/null +++ b/configs/callbacks/curtain_auc_percentiles.yaml @@ -0,0 +1,3 @@ +- _target_: spine.pretrain.curtain.callbacks.CurtainValAUC +- _target_: spine.pretrain.curtain.callbacks.CurtainValLossPercentiles + percentiles: [25, 50, 75] diff --git a/configs/callbacks/curtain_percentiles.yaml b/configs/callbacks/curtain_percentiles.yaml new file mode 100644 index 0000000..5fb5650 --- /dev/null +++ b/configs/callbacks/curtain_percentiles.yaml @@ -0,0 +1,2 @@ +- _target_: spine.pretrain.curtain.callbacks.CurtainValLossPercentiles + percentiles: [25, 50, 75] diff --git a/configs/train.yaml b/configs/train.yaml index 5916748..8cd7278 100644 --- a/configs/train.yaml +++ b/configs/train.yaml @@ -18,6 +18,11 @@ geo: ??? # geometry .npz (required) geo_sensor_key: ??? # per-row sensor-key array in the geometry asset # (e.g. pmt_id) that the reader's sensor_key matches out: ??? # transfer-checkpoint output path (required) +init_from: null # earlier run's TransferCheckpoint (.pth): NEW training + # from its weights, fresh optimizer/scheduler +resume_from: null # Lightning last.ckpt: CONTINUE that same training, full + # state restored; excludes init_from +save_state: null # last.ckpt dir; null = _state/ beside out # Lightning lr_scheduler metadata (monitor/interval/frequency); null uses # SSLModule's default: epoch-level plateau scheduling on the val loss diff --git a/examples/train_curtain.py b/examples/train_curtain.py index 6f11ec6..804103b 100644 --- a/examples/train_curtain.py +++ b/examples/train_curtain.py @@ -83,6 +83,9 @@ def main(cfg: DictConfig) -> None: callbacks=[instantiate(c) for c in (cfg.get("callbacks") or [])], wandb=wandb_cfg, config=OmegaConf.to_container(cfg, resolve=True), + init_from=cfg.get("init_from"), + resume_from=cfg.get("resume_from"), + save_state=cfg.get("save_state"), ) diff --git a/src/spine/pretrain/curtain/callbacks.py b/src/spine/pretrain/curtain/callbacks.py index 15351ba..c8970e9 100644 --- a/src/spine/pretrain/curtain/callbacks.py +++ b/src/spine/pretrain/curtain/callbacks.py @@ -9,6 +9,8 @@ import numpy as np import pytorch_lightning as pl +import torch +import torch.nn.functional as F from pytorch_lightning.callbacks import Callback from spine.pretrain.curtain.task import real_query_mask @@ -121,3 +123,102 @@ def on_validation_epoch_end( pl_module.log("val_auc_hard", auc(lg[hd], y[hd]), sync_dist=True) pl_module.log("val_auc_easy", auc(lg[easy], y[easy]), sync_dist=True) self._cache.clear() + + +class CurtainValLossPercentiles(Callback): + """Percentiles of the per-event validation loss. + + The mean CURTAIN loss is dominated by a few pathological events, so its + curve is noisy and early stopping partly luck; per-event percentiles (of + the mean occupancy BCE over each event's real queries) are stable at the + same cost. Logs `val_loss_p` per configured percentile; monitoring + only, nothing selects on it by default. + """ + + def __init__(self, percentiles: tuple[float, ...] = (25.0, 50.0, 75.0)): + """Configure which percentiles to log. + + Args: + percentiles: Percentiles in [0, 100], each logged as + ``val_loss_p``. + """ + self.percentiles = tuple(percentiles) + self._cache: list = [] + + def on_validation_epoch_start( + self, trainer: pl.Trainer, pl_module: pl.LightningModule + ) -> None: + """Drop caches of a previous epoch. + + Args: + trainer: The running Trainer. + pl_module: The training module. + """ + self._cache.clear() + + def on_validation_batch_end( + self, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: list | None, + batch: dict, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + """Cache this batch's per-event mean occupancy loss. + + Args: + trainer: The running Trainer. + pl_module: The training module; its task locates the occupancy + objective among the outputs. + outputs: validation_step's return -- one prediction tensor per + objective. + batch: The collated batch. + batch_idx: Index of the batch (unused). + dataloader_idx: Index of the dataloader (unused). + """ + occ = next( + ( + i + for i, o in enumerate(pl_module.task.objectives) + if o.name == "occupancy" + ), + None, + ) + if occ is None or outputs is None: + return + pred = outputs[occ] + m = real_query_mask(pred, batch) + per_query = F.binary_cross_entropy_with_logits( + pred[m].squeeze(-1).float(), + batch["label"].values().float(), + reduction="none", + ) + # queries pack per event; the NJT offsets give each event's slice + counts = batch["qpos"].offsets().diff() + ends = torch.cumsum(counts, 0) + starts = ends - counts + per_event = torch.stack( + [per_query[s:e].mean() for s, e in zip(starts, ends, strict=True) if e > s] + ) + self._cache.append(per_event.detach().cpu().numpy()) + + def on_validation_epoch_end( + self, trainer: pl.Trainer, pl_module: pl.LightningModule + ) -> None: + """Log the configured percentiles over this rank's validation events. + + Args: + trainer: The running Trainer. + pl_module: The training module (used for logging). + """ + if not self._cache: + return + v = np.concatenate(self._cache) + # sync_dist averages the ranks' percentiles -- an approximation of + # the global ones, which is fine for a monitoring statistic + for q in self.percentiles: + pl_module.log( + f"val_loss_p{q:g}", float(np.percentile(v, q)), sync_dist=True + ) + self._cache.clear() diff --git a/src/spine/train.py b/src/spine/train.py index d30aab0..934debd 100644 --- a/src/spine/train.py +++ b/src/spine/train.py @@ -6,13 +6,18 @@ from __future__ import annotations +import os from collections.abc import Callable from datetime import timedelta import pytorch_lightning as pl import torch from lightning_fabric.plugins.environments import LightningEnvironment -from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor +from pytorch_lightning.callbacks import ( + EarlyStopping, + LearningRateMonitor, + ModelCheckpoint, +) from pytorch_lightning.strategies import DDPStrategy from torch.utils.data import Dataset @@ -44,6 +49,9 @@ def fit( callbacks: list | None = None, wandb: dict | None = None, config: dict | None = None, + init_from: str | None = None, + resume_from: str | None = None, + save_state: str | None = None, ): """Assemble the datamodule, module and Trainer, then fit. @@ -69,9 +77,19 @@ def fit( wandb: Optional {project, group, name, mode, tags} enabling a WandbLogger + LR monitoring; None trains without a logger. config: Run configuration stored in the checkpoint and logged. + init_from: Earlier run's TransferCheckpoint: start a NEW training + from its weights (fresh optimizer/scheduler). + resume_from: Lightning ``last.ckpt`` to resume from with full state + (optimizer, scheduler, callbacks, loop). Mutually exclusive + with ``init_from``. + save_state: Directory for the rolling full-state ``last.ckpt``, + refreshed each validation epoch; None uses ``_state/``. Returns: The trained SSLModule. + + Raises: + ValueError: If both ``init_from`` and ``resume_from`` are given. """ # fp32 matmuls on TF32 tensor cores: a large speedup on Ampere+ GPUs with # far less precision loss than bf16-mixed @@ -91,12 +109,47 @@ def fit( scheduler=scheduler, scheduler_config=scheduler_config, ) + if init_from is not None and resume_from is not None: + raise ValueError( + "init_from and resume_from are mutually exclusive: a full-state " + "resume already restores the weights" + ) + if init_from is not None: + prior = torch.load(init_from, map_location="cpu", weights_only=False) + module.model.load_state_dict(prior["full_state"]) + print( + f"warm-start: loaded pretext model from {init_from} " + f"(val_loss={prior.get('val_loss')})", + flush=True, + ) + # checked at validation end: a resume replays on_train_epoch_end without + # validation metrics, where the default check would raise cbs = [ TransferCheckpoint(out, config=config or {}), - EarlyStopping(monitor="val_loss_epoch", mode="min", patience=patience), + EarlyStopping( + monitor="val_loss_epoch", + mode="min", + patience=patience, + check_on_train_epoch_end=False, + ), *(callbacks or []), ] + if save_state is None: + save_state = f"{os.path.splitext(out)[0]}_state" + # monitor=None + save_top_k=1: Lightning refreshes last.ckpt only + # alongside a top-k save; save_top_k=0 would defer it to on_train_end, + # useless for crash/timeout recovery + cbs.append( + ModelCheckpoint( + dirpath=save_state, + monitor=None, + save_top_k=1, + save_last=True, + every_n_epochs=1, + save_on_train_epoch_end=False, + ) + ) logger = False if wandb: from pytorch_lightning.loggers import WandbLogger @@ -134,10 +187,10 @@ def fit( max_epochs=max_epochs, gradient_clip_val=grad_clip, num_sanity_val_steps=0, - enable_checkpointing=False, + enable_checkpointing=True, log_every_n_steps=100, logger=logger, callbacks=cbs, ) - trainer.fit(module, datamodule=dm) + trainer.fit(module, datamodule=dm, ckpt_path=resume_from) return module diff --git a/src/spine/utils.py b/src/spine/utils.py index 21d83c7..8696bd7 100644 --- a/src/spine/utils.py +++ b/src/spine/utils.py @@ -30,6 +30,22 @@ def __init__(self, out: str, config: dict, min_delta: float = 1e-4): self.best = float("inf") os.makedirs(os.path.dirname(out) or ".", exist_ok=True) + def state_dict(self) -> dict: + """Persist the export threshold so a full resume stays monotone. + + Returns: + Mapping holding the best exported validation loss. + """ + return {"best": self.best} + + def load_state_dict(self, state_dict: dict) -> None: + """Restore the export threshold on a Lightning resume. + + Args: + state_dict: Mapping produced by :meth:`state_dict`. + """ + self.best = state_dict.get("best", float("inf")) + def on_validation_epoch_end( self, trainer: pl.Trainer, pl_module: pl.LightningModule ) -> None: