feat(dynacell): opt-in channels_last/compile_model flags and LinearLRCooldown callback

.github/workflows/test.yml

@@ -187,6 +187,7 @@ jobs:
             ../../packages/viscy-utils/tests/test_prediction_writer.py \
             ../../packages/viscy-utils/tests/test_prediction_metadata.py \
             ../../packages/viscy-utils/src/viscy_utils/callbacks/monitor_health_test.py \
+            ../../packages/viscy-utils/src/viscy_utils/callbacks/linear_lr_cooldown_test.py \
             -v
         working-directory: applications/dynacell
 

applications/dynacell/src/dynacell/engine.py

@@ -16,6 +16,7 @@
 import torch
 import torch.nn.functional as F
 from lightning.pytorch import Callback, LightningModule
+from lightning.pytorch.trainer.states import TrainerFn
 from monai.transforms import DivisiblePad
 from torch import Tensor, nn
 
@@ -592,6 +593,24 @@ class DynacellUNet(LightningModule):
         must then carry ``fg_mask`` (set ``fg_mask_key`` on the datamodule).
     seg_aux_weight : float
         Weight of the ``seg_aux`` term. Must be ``0.0`` when ``seg_aux`` is None.
+    channels_last : bool, default False
+        Store the backbone's convolution weights in channels-last layout
+        (``torch.channels_last_3d`` for 5-D weights, ``torch.channels_last``
+        for 4-D ones, since UNeXt2 mixes 2-D and 3-D layers) and feed 5-D
+        inputs as ``channels_last_3d``. Only strides change, so outputs and
+        state-dict keys and values match the default layout. Applied in
+        :meth:`configure_model`, before the strategy moves the module to the
+        device and wraps it in DDP; converting after the DDP wrap breaks its
+        gradient bucket views. Not saved in the hyperparameters, so a resumed
+        fit takes the value from its own config.
+    compile_model : bool, default False
+        During fit, run the backbone through ``torch.compile(self.model)``
+        kept in a plain-dict cache, not as a registered submodule, so the
+        state-dict keys never change. Every other stage (validate, predict)
+        and any later use of the same module run the eager backbone:
+        sliding-window predict varies the input shape. Built in
+        :meth:`configure_model`, before the DDP wrap, and dropped when fit
+        ends or raises. Not saved in the hyperparameters.
     """
 
     def __init__(
@@ -614,9 +633,16 @@ def __init__(
         encoder_only: bool = False,
         seg_aux: SegAuxDice | None = None,
         seg_aux_weight: float = 0.0,
+        channels_last: bool = False,
+        compile_model: bool = False,
     ) -> None:
         super().__init__()
-        self.save_hyperparameters(ignore=["loss_function", "ckpt_path", "encoder_only", "seg_aux"])
+        # channels_last and compile_model change speed, not the model, so a
+        # resume takes them from its own config and older checkouts can still
+        # resume checkpoints written with them on.
+        self.save_hyperparameters(
+            ignore=["loss_function", "ckpt_path", "encoder_only", "seg_aux", "channels_last", "compile_model"]
+        )
         if model_config is None:
             model_config = {}
         net_class = _ARCHITECTURE.get(architecture)
@@ -636,6 +662,10 @@ def __init__(
         self.predict_phase_shifts = predict_phase_shifts
         self.seg_aux = seg_aux
         self.seg_aux_weight = seg_aux_weight
+        self.channels_last = channels_last
+        self.compile_model = compile_model
+        # Fit-scoped compiled backbone; a plain dict keeps it out of the module tree.
+        self._compiled: dict[str, Callable[[Tensor], Tensor]] = {}
 
         self.training_step_outputs: list = []
         # Each entry is a list of (loss, batch_size) tuples for weighted aggregation.
@@ -675,20 +705,51 @@ def __init__(
         elif ckpt_path is not None:
             self.load_state_dict(_ckpt_state_dict(ckpt_path))
 
+    def configure_model(self) -> None:
+        """Apply the opt-in ``channels_last`` and ``compile_model`` settings.
+
+        Lightning calls this before the strategy moves the module to the device
+        and wraps it in DDP; both settings must precede the wrap. The hook may
+        run more than once, so both steps are idempotent.
+        """
+        if self.channels_last:
+            for param in self.model.parameters():
+                if param.ndim == 5:
+                    param.data = param.data.contiguous(memory_format=torch.channels_last_3d)
+                elif param.ndim == 4:
+                    param.data = param.data.contiguous(memory_format=torch.channels_last)
+        if self.compile_model and self.trainer.state.fn == TrainerFn.FITTING and "fit" not in self._compiled:
+            self._compiled["fit"] = torch.compile(self.model)
+
+    def drop_compiled_model(self) -> None:
+        """Drop the fit-scoped compiled backbone, so later stages run the eager one."""
+        self._compiled.clear()
+
+    def configure_callbacks(self) -> list[Callback]:
+        """Attach the callback that drops the compiled backbone when a fit raises."""
+        return [_DropCompiledModelOnException()] if self.compile_model else []
+
+    def on_fit_end(self) -> None:
+        """Drop the fit-scoped compiled backbone."""
+        self.drop_compiled_model()
+
     def forward(self, x: Tensor) -> Tensor:
         """Run forward pass through the model.
 
         Parameters
         ----------
         x : Tensor
-            Input tensor of shape ``(B, C, D, H, W)``.
+            Input tensor of shape ``(B, C, D, H, W)``. Made
+            ``channels_last_3d`` when ``channels_last`` is set.
 
         Returns
         -------
         Tensor
             Model output.
         """
-        return self.model(x)
+        if self.channels_last and x.ndim == 5:
+            x = x.contiguous(memory_format=torch.channels_last_3d)
+        return self._compiled.get("fit", self.model)(x)
 
     def _compute_loss(self, pred: Tensor, target: Tensor, batch: Sample) -> tuple[Tensor, dict[str, Tensor]]:
         """Compute loss, optionally passing fg_mask to the loss function.
@@ -918,6 +979,18 @@ def predict_sliding_window(self, source: Tensor, overlap_size: tuple[int, int, i
         )
 
 
+class _DropCompiledModelOnException(Callback):
+    """Drop :class:`DynacellUNet`'s fit-scoped compiled backbone when a Trainer run raises.
+
+    Lightning skips ``on_fit_end`` on an exception, so without this a later stage
+    on the same module would still run the compiled callable.
+    """
+
+    def on_exception(self, trainer, pl_module: LightningModule, exception: BaseException) -> None:
+        """Drop the compiled backbone."""
+        pl_module.drop_compiled_model()
+
+
 class _DropCompiledNetOnException(Callback):
     """Free :class:`DynacellFlowMatching`'s compiled inference copy when a Trainer run raises.
 

applications/dynacell/tests/test_engine.py

@@ -1902,3 +1902,156 @@ def test_flow_matching_generate_trajectory_fixed_grid(synth_celldiff_batch):
     # Midpoint takes 2 evaluations per step on the 8-point grid: 7 steps.
     assert len(calls) == 14
     torch.testing.assert_close(traj[-1], final)
+
+
+# ---------------------------------------------------------------------------
+# Opt-in throughput flags: channels_last and compile_model
+# ---------------------------------------------------------------------------
+
+
+class _PairDataset(torch.utils.data.Dataset):
+    """Fixed-shape (source, target) Sample dicts for short CPU fits."""
+
+    def __init__(self, n: int, shape: tuple[int, int, int, int]):
+        generator = torch.Generator().manual_seed(0)
+        self.source = torch.randn(n, *shape, generator=generator)
+        self.target = torch.randn(n, *shape, generator=generator)
+
+    def __len__(self) -> int:
+        return len(self.source)
+
+    def __getitem__(self, idx: int) -> dict:
+        return {"source": self.source[idx], "target": self.target[idx]}
+
+
+@pytest.mark.parametrize(
+    ("architecture", "config", "shape"),
+    [("FNet3D", FNET_TEST_CONFIG, (2, 1, 4, 16, 16)), ("UNeXt2", UNEXT2_TEST_CONFIG, (2, 1, 5, 64, 64))],
+)
+def test_dynacell_unet_channels_last_matches_default_layout(architecture, config, shape):
+    """channels_last changes strides only: same weights give the same output."""
+    seed_everything(0)
+    default = DynacellUNet(architecture=architecture, model_config=config).eval()
+    fast = DynacellUNet(architecture=architecture, model_config=config, channels_last=True).eval()
+    fast.load_state_dict(default.state_dict())
+    fast.configure_model()
+    fast.configure_model()  # Lightning may call the hook more than once.
+
+    # Only weights whose default layout is not already channels-last can show
+    # the conversion (C_in == 1 or 1x1 kernels satisfy both layouts).
+    formats = {5: torch.channels_last_3d, 4: torch.channels_last}
+    converted = {ndim: 0 for ndim in formats}
+    for before, after in zip(default.model.parameters(), fast.model.parameters(), strict=True):
+        if before.ndim in formats and not before.is_contiguous(memory_format=formats[before.ndim]):
+            assert after.is_contiguous(memory_format=formats[after.ndim])
+            converted[before.ndim] += 1
+    assert converted[5] > 0
+    if architecture == "UNeXt2":
+        # UNeXt2 mixes 2-D and 3-D layers; its 4-D weights go channels_last.
+        assert converted[4] > 0
+
+    x = torch.randn(shape)
+    with torch.no_grad():
+        expected = default(x)
+        actual = fast(x)
+    assert torch.allclose(actual, expected, atol=1e-5, rtol=1e-4)
+
+
+def test_dynacell_unet_throughput_flags_keep_checkpoints_interchangeable(tmp_path):
+    """A fit with both flags on writes a checkpoint a flags-off model loads unchanged.
+
+    The flags are not saved in the hyperparameters, the compiled backbone is not
+    a registered submodule (no ``_orig_mod.`` keys), and the channels-last weights
+    load into the default layout with identical values.
+    """
+    seed_everything(0)
+    module = DynacellUNet(
+        architecture="FNet3D",
+        model_config=FNET_TEST_CONFIG,
+        channels_last=True,
+        compile_model=True,
+        log_batches_per_epoch=0,
+    )
+    assert "channels_last" not in module.hparams
+    assert "compile_model" not in module.hparams
+    reference_keys = list(DynacellUNet(architecture="FNet3D", model_config=FNET_TEST_CONFIG).state_dict())
+
+    trainer = Trainer(
+        max_steps=1,
+        accelerator="cpu",
+        devices=1,
+        logger=False,
+        enable_checkpointing=False,
+        enable_progress_bar=False,
+        enable_model_summary=False,
+        limit_val_batches=0,
+    )
+    compiled_during_fit: list[bool] = []
+
+    class _RecordCompiled(Callback):
+        def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+            compiled_during_fit.append("fit" in pl_module._compiled)
+
+    trainer.callbacks.append(_RecordCompiled())
+    loader = torch.utils.data.DataLoader(_PairDataset(2, (1, 4, 16, 16)), batch_size=2)
+    trainer.fit(module, train_dataloaders=loader)
+    assert compiled_during_fit == [True]
+    assert not module._compiled  # dropped when fit ends
+    assert list(module.state_dict()) == reference_keys
+
+    ckpt_path = tmp_path / "flags_on.ckpt"
+    trainer.save_checkpoint(ckpt_path)
+    saved = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    assert "channels_last" not in saved["hyper_parameters"]
+    assert "compile_model" not in saved["hyper_parameters"]
+
+    restored = DynacellUNet(architecture="FNet3D", model_config=FNET_TEST_CONFIG, ckpt_path=str(ckpt_path))
+    trained = module.state_dict()
+    for key, value in restored.state_dict().items():
+        assert torch.equal(value, trained[key]), key
+    reloaded = DynacellUNet.load_from_checkpoint(ckpt_path, model_config=FNET_TEST_CONFIG)
+    assert reloaded.channels_last is False and reloaded.compile_model is False
+
+
+def test_dynacell_unet_compile_model_is_fit_scoped(monkeypatch):
+    """The compiled backbone exists only during fit: predict, alone or after a fit or a
+    failed fit on the same module, runs the eager backbone."""
+    monkeypatch.setattr(torch, "compile", lambda module, **kwargs: lambda x: module(x))  # fast stand-in
+    module = DynacellUNet(
+        architecture="FNet3D", model_config=FNET_TEST_CONFIG, compile_model=True, log_batches_per_epoch=0
+    )
+    loader = torch.utils.data.DataLoader(_PairDataset(2, (1, 4, 16, 16)), batch_size=2)
+    compiled_at: list[tuple[str, bool]] = []
+
+    class _Record(Callback):
+        def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+            compiled_at.append(("fit", "fit" in pl_module._compiled))
+
+        def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
+            compiled_at.append(("predict", "fit" in pl_module._compiled))
+
+    class _Fail(Callback):
+        def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+            raise RuntimeError("injected fit failure")
+
+    def trainer(*callbacks: Callback) -> Trainer:
+        return Trainer(
+            max_steps=1,
+            accelerator="cpu",
+            devices=1,
+            logger=False,
+            enable_checkpointing=False,
+            enable_progress_bar=False,
+            enable_model_summary=False,
+            limit_val_batches=0,
+            callbacks=[_Record(), *callbacks],
+        )
+
+    trainer().predict(module, dataloaders=loader)
+    trainer().fit(module, train_dataloaders=loader)
+    trainer().predict(module, dataloaders=loader)
+    with pytest.raises(RuntimeError, match="injected fit failure"):
+        trainer(_Fail()).fit(module, train_dataloaders=loader)
+    assert not module._compiled
+    trainer().predict(module, dataloaders=loader)
+    assert compiled_at == [("predict", False), ("fit", True), ("predict", False), ("fit", True), ("predict", False)]

packages/viscy-utils/src/viscy_utils/callbacks/__init__.py

@@ -1,5 +1,6 @@
 from viscy_utils.callbacks.embedding_snapshot import EmbeddingSnapshotCallback
 from viscy_utils.callbacks.embedding_writer import EmbeddingWriter
+from viscy_utils.callbacks.linear_lr_cooldown import LinearLRCooldown
 from viscy_utils.callbacks.monitor_health import MonitorHealthCheck, MonitorHealthError
 from viscy_utils.callbacks.online_eval import OnlineEvalCallback
 from viscy_utils.callbacks.optimizer_health import OptimizerHealthCheck, OptimizerHealthError
@@ -8,6 +9,7 @@
 __all__ = [
     "EmbeddingSnapshotCallback",
     "EmbeddingWriter",
+    "LinearLRCooldown",
     "MonitorHealthCheck",
     "MonitorHealthError",
     "OnlineEvalCallback",