Avoid mutating RGB tensors in perceptual normalization

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:

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()