Lossless stop/resume for fit(): rolling full-state last.ckpt + resume_from, init_from warm start

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]

configs/callbacks/curtain_percentiles.yaml

@@ -0,0 +1,2 @@
+- _target_: spine.pretrain.curtain.callbacks.CurtainValLossPercentiles
+  percentiles: [25, 50, 75]

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 = <out stem>_state/ beside out
 
 # Lightning lr_scheduler metadata (monitor/interval/frequency); null uses
 # SSLModule's default: epoch-level plateau scheduling on the val loss

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"),
     )
 
 

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<q>` 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<q>``.
+        """
+        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()

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 ``<out stem>_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

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: