Prevent empty slice sampling in fake-3D perceptual loss

monai/losses/perceptual.py

@@ -66,6 +66,7 @@ class PerceptualLoss(nn.Module):
             - "resnet50"
         is_fake_3d: if True use 2.5D approach for a 3D perceptual loss.
         fake_3d_ratio: ratio of how many slices per axis are used in the 2.5D approach.
+            At least one slice per axis is used.
         cache_dir: path to cache directory to save the pretrained network weights.
         pretrained: whether to load pretrained weights. This argument only works when using networks from
             LIPIS or Torchvision. Defaults to ``True``.
@@ -183,9 +184,8 @@ def batchify_axis(x: torch.Tensor, fake_3d_perm: tuple) -> torch.Tensor:
 
         channel_axis = 1
         input_slices = batchify_axis(x=input, fake_3d_perm=(spatial_axis, channel_axis) + tuple(preserved_axes))
-        indices = torch.randperm(input_slices.shape[0])[: int(input_slices.shape[0] * self.fake_3d_ratio)].to(
-            input_slices.device
-        )
+        num_slices = max(1, int(input_slices.shape[0] * self.fake_3d_ratio))
+        indices = torch.randperm(input_slices.shape[0])[:num_slices].to(input_slices.device)
         input_slices = torch.index_select(input_slices, dim=0, index=indices)
         target_slices = batchify_axis(x=target, fake_3d_perm=(spatial_axis, channel_axis) + tuple(preserved_axes))
         target_slices = torch.index_select(target_slices, dim=0, index=indices)

tests/losses/test_perceptual_loss.py

@@ -100,6 +100,17 @@ def test_identical_input(self, input_param, input_shape, target_shape):
         else:
             self.assertEqual(result, torch.Tensor([0.0]))
 
+    @parameterized.expand([0.01, 0.2, 1 / 3, 1.0])
+    def test_fake_3d_small_ratio(self, fake_3d_ratio):
+        loss = PerceptualLoss(spatial_dims=3, network_type="resnet50", pretrained=False, fake_3d_ratio=fake_3d_ratio)
+        tensor = torch.rand(1, 1, 3, 32, 32, requires_grad=True)
+        result = loss(tensor, tensor)
+
+        self.assertEqual(result.item(), 0.0)
+        result.backward()
+        self.assertIsNotNone(tensor.grad)
+        self.assertTrue(torch.isfinite(tensor.grad).all())
+
     def test_different_shape(self):
         with skip_if_downloading_fails():
             loss = PerceptualLoss(spatial_dims=2, network_type="squeeze")