diff --git a/monai/losses/perceptual.py b/monai/losses/perceptual.py index 6c0608f605..d412c44244 100644 --- a/monai/losses/perceptual.py +++ b/monai/losses/perceptual.py @@ -468,12 +468,9 @@ def spatial_average(x: torch.Tensor, keepdim: bool = True) -> torch.Tensor: def torchvision_zscore_norm(x: torch.Tensor) -> torch.Tensor: - mean = [0.485, 0.456, 0.406] - std = [0.229, 0.224, 0.225] - x[:, 0, :, :] = (x[:, 0, :, :] - mean[0]) / std[0] - x[:, 1, :, :] = (x[:, 1, :, :] - mean[1]) / std[1] - x[:, 2, :, :] = (x[:, 2, :, :] - mean[2]) / std[2] - return x + mean = x.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = x.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + return (x - mean) / std def subtract_mean(x: torch.Tensor) -> torch.Tensor: diff --git a/tests/losses/test_perceptual_loss.py b/tests/losses/test_perceptual_loss.py index 051a79fef1..24c6b605b9 100644 --- a/tests/losses/test_perceptual_loss.py +++ b/tests/losses/test_perceptual_loss.py @@ -17,7 +17,7 @@ from parameterized import parameterized from monai.losses import PerceptualLoss -from monai.losses.perceptual import normalize_tensor +from monai.losses.perceptual import normalize_tensor, torchvision_zscore_norm from monai.utils import optional_import from tests.test_utils import assert_allclose, skip_if_downloading_fails, skip_if_quick @@ -138,5 +138,57 @@ def test_normalize_tensor_zero_norm_finite_gradient(self): self.assertFalse(torch.isnan(x.grad).any()) +@unittest.skipUnless(has_torchvision, "Requires torchvision") +class TestTorchvisionPerceptualLoss(unittest.TestCase): + @parameterized.expand([1, 3]) + def test_input_not_modified(self, channels): + loss = PerceptualLoss(spatial_dims=2, network_type="resnet50", pretrained=False) + input = torch.rand(1, channels, 32, 32) + target = torch.rand_like(input) + input_before = input.clone() + target_before = target.clone() + + result = loss(input, target) + + self.assertTrue(torch.isfinite(result)) + self.assertTrue(torch.equal(input, input_before)) + self.assertTrue(torch.equal(target, target_before)) + + @parameterized.expand([1, 3]) + def test_leaf_gradient(self, channels): + loss = PerceptualLoss(spatial_dims=2, network_type="resnet50", pretrained=False) + input = torch.rand(1, channels, 32, 32, requires_grad=True) + target = torch.rand_like(input, requires_grad=True) + + result = loss(input, target) + result.backward() + + self.assertIsNotNone(input.grad) + self.assertIsNotNone(target.grad) + self.assertTrue(torch.isfinite(input.grad).all()) + self.assertTrue(torch.isfinite(target.grad).all()) + + +class TestTorchvisionZScoreNorm(unittest.TestCase): + @parameterized.expand([torch.float32, torch.float64]) + def test_normalization(self, dtype): + input = torch.rand(2, 3, 4, 4, dtype=dtype) + input_before = input.clone() + expected = torch.stack( + [ + (input[:, channel] - mean) / std + for channel, (mean, std) in enumerate(zip([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])) + ], + dim=1, + ) + + result = torchvision_zscore_norm(input) + + assert_allclose(result, expected, rtol=0, atol=0) + self.assertEqual(result.dtype, input.dtype) + self.assertEqual(result.device, input.device) + self.assertTrue(torch.equal(input, input_before)) + + if __name__ == "__main__": unittest.main()