From 5d62bc6be8d2e00933d12e338d4ff2c83f16e754 Mon Sep 17 00:00:00 2001 From: Don Yin Date: Tue, 20 Oct 2026 19:44:58 +0100 Subject: [PATCH] Keep fake-3D perceptual loss slice sampling nonempty Clamp the selected slice count to one for nonempty inputs when a small sampling ratio would otherwise produce an empty mean. Add regression coverage for small ratios, finite gradients, and the unchanged controls. Signed-off-by: Don Yin Assisted-by: Claude --- monai/losses/perceptual.py | 6 +++--- tests/losses/test_perceptual_loss.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/monai/losses/perceptual.py b/monai/losses/perceptual.py index 6c0608f605a..4b90771af61 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 051a79fef13..d68b71b6b32 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")