fix NaN gradients in 3d triangulation loss (#358)

Author: themattinthehatt

lightning_pose/losses/losses.py

@@ -776,17 +776,55 @@ def remove_nans(
         loss: TensorType["batch", "cam_pairs", "num_keypoints"],
     ) -> TensorType["valid_losses"]:
         mask = ~torch.isnan(loss)
-        if mask.sum() == 0.0:
-            return torch.tensor(0.0, device=loss.device, dtype=loss.dtype)
+        valid_losses = torch.masked_select(loss, mask)
+        if valid_losses.numel() == 0:
+            # No valid losses, return zero that preserves gradients
+            # Use torch.where to avoid nan*0.0 issues
+            dummy_loss = torch.where(mask, loss, torch.zeros_like(loss))
+            return dummy_loss.sum()  # This will be 0.0 and preserve gradients
         else:
-            return torch.masked_select(loss, ~torch.isnan(loss))
+            return valid_losses
 
     def compute_loss(
         self,
         targets: TensorType["batch", "num_keypoints", 3],
         predictions: TensorType["batch", "cam_pairs", "num_keypoints", 3],
     ) -> TensorType["batch", "cam_pairs", "num_keypoints"]:
-        loss = torch.linalg.norm(targets.unsqueeze(1) - predictions, ord=2, dim=-1)
+
+        # Check for NaN targets AND predictions
+        nan_targets = torch.isnan(targets).any(dim=-1)  # [batch, num_keypoints]
+        nan_predictions = torch.isnan(predictions).any(dim=-1)  # [batch, cam_pairs, num_keypoints]
+
+        # Expand target NaN mask to match prediction dimensions
+        nan_targets_expanded = nan_targets.unsqueeze(1)  # [batch, 1, num_keypoints]
+
+        # Combined NaN mask
+        combined_nan_mask = \
+            nan_targets_expanded | nan_predictions  # [batch, cam_pairs, num_keypoints]
+
+        # Create clean targets and predictions - replace NaNs with zeros and detach
+        clean_targets = torch.where(
+            nan_targets.unsqueeze(-1),  # [batch, num_keypoints, 1]
+            torch.zeros_like(targets).detach(),
+            targets,
+        )
+
+        clean_predictions = torch.where(
+            combined_nan_mask.unsqueeze(-1),  # [batch, cam_pairs, num_keypoints, 1]
+            torch.zeros_like(predictions).detach(),
+            predictions,
+        )
+
+        # Compute loss with clean tensors
+        loss = torch.linalg.norm(clean_targets.unsqueeze(1) - clean_predictions, ord=2, dim=-1)
+
+        # Set loss to NaN where either targets or predictions were originally NaN
+        loss = torch.where(
+            combined_nan_mask,
+            torch.tensor(float('nan'), device=loss.device, dtype=loss.dtype),
+            loss,
+        )
+
         return loss
 
     def __call__(

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "lightning-pose"
-version = "2.0.3"
+version = "2.0.4"
 description = "Semi-supervised pose estimation using pytorch lightning"
 license = "MIT"
 readme = "README.md"

tests/losses/test_losses.py

@@ -508,30 +508,45 @@ def test_targets_all_nans(self, pp_loss):
         num_keypoints = 4
         num_cam_pairs = 3
         keypoints_targ_3d = torch.full((num_batch, num_keypoints, 3), float('nan'))
-        keypoints_pred_3d = torch.ones((num_batch, num_cam_pairs, num_keypoints, 3))
+        keypoints_pred_3d = torch.ones(
+            (num_batch, num_cam_pairs, num_keypoints, 3),
+            requires_grad=True,
+        )
         loss, _ = pp_loss(keypoints_targ_3d, keypoints_pred_3d)
         assert loss.item() == 0.0
+        loss.backward()
+        assert not torch.isnan(keypoints_pred_3d.grad).any(), "gradients contain NaN values"
 
     def test_predictions_all_nans(self, pp_loss):
         num_batch = 1
         num_keypoints = 4
         num_cam_pairs = 3
         keypoints_targ_3d = torch.ones((num_batch, num_keypoints, 3))
-        keypoints_pred_3d = torch.full((num_batch, num_cam_pairs, num_keypoints, 3), float('nan'))
+        keypoints_pred_3d = torch.full(
+            (num_batch, num_cam_pairs, num_keypoints, 3), float('nan'),
+            requires_grad=True,
+        )
         loss, _ = pp_loss(keypoints_targ_3d, keypoints_pred_3d)
         assert loss.item() == 0.0
+        loss.backward()
+        assert not torch.isnan(keypoints_pred_3d.grad).any(), "gradients contain NaN values"
 
     def test_targets_partial_nans(self, pp_loss):
         num_batch = 2
         num_keypoints = 4
         num_cam_pairs = 2
         keypoints_targ_3d = torch.zeros(size=(num_batch, num_keypoints, 3))
         keypoints_targ_3d[0, 0, :] = float('nan')  # first keypoint in first batch NaN
-        keypoints_pred_3d = torch.ones(size=(num_batch, num_cam_pairs, num_keypoints, 3))
+        keypoints_pred_3d = torch.ones(
+            size=(num_batch, num_cam_pairs, num_keypoints, 3),
+            requires_grad=True,
+        )
         loss, _ = pp_loss(keypoints_targ_3d, keypoints_pred_3d)
         # each valid position has loss = sqrt(3) (distance from 0 to 1 in 3D)
         expected_loss = torch.sqrt(torch.tensor(3.0))
         assert loss.isclose(expected_loss)
+        loss.backward()
+        assert not torch.isnan(keypoints_pred_3d.grad).any(), "gradients contain NaN values"
 
     def test_predictions_partial_nans(self, pp_loss):
         num_batch = 3
@@ -542,10 +557,13 @@ def test_predictions_partial_nans(self, pp_loss):
         keypoints_pred_3d[0, 0, 0, :] = float('nan')
         keypoints_pred_3d[1, 1, :, :] = float('nan')
         keypoints_pred_3d[2, :, :, :] = float('nan')
+        keypoints_pred_3d.requires_grad_(True)  # need to do this after inplace operations
         loss, _ = pp_loss(keypoints_targ_3d, keypoints_pred_3d)
         # each valid position has loss = sqrt(3) (distance from 0 to 1 in 3D)
         expected_loss = torch.sqrt(torch.tensor(3.0))
         assert loss.isclose(expected_loss)
+        loss.backward()
+        assert not torch.isnan(keypoints_pred_3d.grad).any(), "gradients contain NaN values"
 
 
 def test_get_loss_classes():