diff --git a/src/monai_physio/train_physicsnemo_physics_informed_motion.py b/src/monai_physio/train_physicsnemo_physics_informed_motion.py index db8b397..9a5b25f 100644 --- a/src/monai_physio/train_physicsnemo_physics_informed_motion.py +++ b/src/monai_physio/train_physicsnemo_physics_informed_motion.py @@ -522,6 +522,9 @@ def __init__(self, log_level: int | str = logging.INFO) -> None: self._epoch_data_loss: Optional[torch.Tensor] = None self._epoch_physics_loss: Optional[torch.Tensor] = None self._epoch_batches = 0 + # Consecutive epochs with physics_mean ~ 0 while lambda_physics > 0; + # see the collapse warning in _log_epoch. + self._zero_physics_streak = 0 def set_mechanics( self, @@ -833,6 +836,21 @@ def _log_epoch(self, context: DistributedContext, epoch: int, epochs: int) -> No divisor = max(batches, 1) physics_mean = physics_sum / divisor + + if self.lambda_physics > 0.0 and physics_mean < 1e-9: + self._zero_physics_streak += 1 + else: + self._zero_physics_streak = 0 + if self._zero_physics_streak == 5 and context.is_main: + self.log_warning( + "physics loss has been ~0 for 5 consecutive epochs while " + "lambda_physics=%.4g; the network likely collapsed to the " + "trivial uniform-displacement solution described in " + "set_lambda_physics_warmup()'s docstring. Consider a longer " + "warmup or a lower lambda_physics target.", + self.lambda_physics, + ) + self._log_main( context, # physics/weighted in scientific notation: %f rounds anything diff --git a/tests/test_physics_informed_motion.py b/tests/test_physics_informed_motion.py index d7e2d1f..283df2f 100644 --- a/tests/test_physics_informed_motion.py +++ b/tests/test_physics_informed_motion.py @@ -502,6 +502,56 @@ def test_the_epoch_log_separates_the_two_loss_terms() -> None: ) +def test_zero_physics_streak_warns_once_after_five_epochs() -> None: + """A collapsed physics term should be flagged, not silently trained through. + + ``set_lambda_physics_warmup``'s docstring describes a trivial-solution + collapse where the physics term reads exactly zero for the rest of the + run. Five consecutive zero-physics epochs should raise exactly one + warning -- not one per epoch, and not at all before the streak or once + physics recovers. + """ + import torch + + from monai_physio.process_physicsnemo import DistributedContext + from monai_physio.train_physicsnemo_physics_informed_motion import ( + TrainPhysicsNeMoPhysicsInformedMotion, + ) + + method = TrainPhysicsNeMoPhysicsInformedMotion() + method.lambda_physics = 0.1 + + warnings: list[str] = [] + method.log_warning = lambda *args: warnings.append(str(args[0]) % args[1:]) # type: ignore[method-assign] + method.log_info = lambda *args: None # type: ignore[method-assign] + + context = DistributedContext( + device=torch.device("cpu"), rank=0, local_rank=0, world_size=1 + ) + + def log_epoch_with(physics: float, epoch: int) -> None: + method._epoch_data_loss = torch.tensor(1.0) + method._epoch_physics_loss = torch.tensor(physics) + method._epoch_batches = 1 + method._log_epoch(context, epoch=epoch, epochs=10) + + for epoch in range(4): + log_epoch_with(0.0, epoch) + assert not warnings, "Should not warn before five consecutive zero epochs" + + log_epoch_with(0.0, 4) + assert len(warnings) == 1, "Should warn exactly once at the fifth zero epoch" + assert "collapsed" in warnings[0] + + log_epoch_with(0.0, 5) + assert len(warnings) == 1, "Should not warn again every epoch after" + + log_epoch_with(1.0, 6) + for epoch in range(4): + log_epoch_with(0.0, 7 + epoch) + assert len(warnings) == 1, "A streak broken by recovery should not re-warn early" + + def test_bind_reference_meshes_repairs_against_template_elements(tmp_path: Any) -> None: """Repair must use ``self._tets``, not whatever cells the file stores.