diff --git a/monai/losses/perceptual.py b/monai/losses/perceptual.py index 6c0608f605..4b90771af6 100644 --- a/monai/losses/perceptual.py +++ b/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) diff --git a/tests/losses/test_perceptual_loss.py b/tests/losses/test_perceptual_loss.py index 051a79fef1..d68b71b6b3 100644 --- a/tests/losses/test_perceptual_loss.py +++ b/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")