feat: improve how device switch is handled between the metric device and the input tensors device

ignite/metrics/ssim.py

@@ -1,4 +1,5 @@
-from typing import Callable, Sequence, Union
+import warnings
+from typing import Callable, Optional, Sequence, Union
 
 import torch
 import torch.nn.functional as F
@@ -102,7 +103,8 @@
         self.c2 = (k2 * data_range) ** 2
         self.pad_h = (self.kernel_size[0] - 1) // 2
         self.pad_w = (self.kernel_size[1] - 1) // 2
-        self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)
+        self._kernel_2d = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)
+        self._kernel: Optional[torch.Tensor] = None
 
     @reinit__is_reduced
     def reset(self) -> None:
@@ -155,9 +157,22 @@
                 f"Expected y_pred and y to have BxCxHxW shape. Got y_pred: {y_pred.shape} and y: {y.shape}."
             )
 
-        channel = y_pred.size(1)
-        if len(self._kernel.shape) < 4:
-            self._kernel = self._kernel.expand(channel, 1, -1, -1).to(device=y_pred.device)
+        nb_channel = y_pred.size(1)
+        if self._kernel is None or self._kernel.shape[0] != nb_channel:
+            self._kernel = self._kernel_2d.expand(nb_channel, 1, -1, -1)
+
+        if y_pred.device != self._kernel.device:
+            if self._kernel.device == torch.device("cpu"):
+                self._kernel = self._kernel.to(device=y_pred.device)
+
+            elif y_pred.device == torch.device("cpu"):
+                warnings.warn(
+                    "y_pred tensor is on cpu device but previous computation was on another device: "
+                    f"{self._kernel.device}. To avoid having a performance hit, please ensure that all "
+                    "y and y_pred tensors are on the same device.",
+                )
+                y_pred = y_pred.to(device=self._kernel.device)
+                y = y.to(device=self._kernel.device)
 
         y_pred = F.pad(y_pred, [self.pad_w, self.pad_w, self.pad_h, self.pad_h], mode="reflect")
         y = F.pad(y, [self.pad_w, self.pad_w, self.pad_h, self.pad_h], mode="reflect")
@@ -166,7 +181,7 @@
             self._kernel = self._kernel.to(dtype=y_pred.dtype)
 
         input_list = [y_pred, y, y_pred * y_pred, y * y, y_pred * y]
-        outputs = F.conv2d(torch.cat(input_list), self._kernel, groups=channel)
+        outputs = F.conv2d(torch.cat(input_list), self._kernel, groups=nb_channel)
         batch_size = y_pred.size(0)
         output_list = [outputs[x * batch_size : (x + 1) * batch_size] for x in range(len(input_list))]
 
@@ -184,7 +199,7 @@
         b2 = sigma_pred_sq + sigma_target_sq + self.c2
 
         ssim_idx = (a1 * a2) / (b1 * b2)
-        self._sum_of_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64).sum().to(self._device)
+        self._sum_of_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64).sum().to(device=self._device)
 
         self._num_examples += y.shape[0]
 

tests/ignite/metrics/test_ssim.py

@@ -1,3 +1,5 @@
+from typing import Sequence, Union
+
 import numpy as np
 import pytest
 import torch
@@ -70,25 +72,49 @@ def test_invalid_ssim():
     "shape, kernel_size, gaussian, use_sample_covariance",
     [[(8, 3, 224, 224), 7, False, True], [(12, 3, 28, 28), 11, True, False]],
 )
-def test_ssim(
-    available_device, shape, kernel_size, gaussian, use_sample_covariance, dtype=torch.float32, precision=7e-5
-):
-    y_pred = torch.rand(shape, device=available_device, dtype=dtype)
+def test_ssim(available_device, shape, kernel_size, gaussian, use_sample_covariance):
+    y_pred = torch.rand(shape, device=available_device)
     y = y_pred * 0.8
 
+    compare_ssim_ignite_skiimg(
+        y_pred,
+        y,
+        available_device,
+        kernel_size=kernel_size,
+        gaussian=gaussian,
+        use_sample_covariance=use_sample_covariance,
+    )
+
+
+def compare_ssim_ignite_skiimg(
+    y_pred: torch.Tensor,
+    y: torch.Tensor,
+    device: torch.device,
+    precision: float = 2e-5,  # default to float32 expected precision
+    *,
+    skimg_y_pred: Union[np.ndarray, None] = None,
+    skimg_y: Union[np.ndarray, None] = None,
+    data_range: float = 1.0,
+    kernel_size: Union[int, Sequence[int]] = 11,
+    gaussian: bool = True,
+    use_sample_covariance: bool = False,
+):
     sigma = 1.5
-    data_range = 1.0
-    ssim = SSIM(data_range=data_range, sigma=sigma, device=available_device)
+
+    ssim = SSIM(data_range=data_range, sigma=sigma, device=device)
     ssim.update((y_pred, y))
     ignite_ssim = ssim.compute()
 
     if y_pred.dtype == torch.bfloat16:
         y_pred = y_pred.to(dtype=torch.float16)
 
-    skimg_pred = y_pred.cpu().numpy()
-    skimg_y = skimg_pred * 0.8
+    if skimg_y_pred is None:
+        skimg_y_pred = y_pred.cpu().numpy()
+    if skimg_y is None:
+        skimg_y = skimg_y_pred * 0.8
+
     skimg_ssim = ski_ssim(
-        skimg_pred,
+        skimg_y_pred,
         skimg_y,
         win_size=kernel_size,
         sigma=sigma,
@@ -102,6 +128,43 @@ def test_ssim(
     assert np.allclose(ignite_ssim, skimg_ssim, atol=precision)
 
 
+@pytest.mark.parametrize(
+    "metric_device, y_pred_device",
+    [
+        [torch.device("cpu"), torch.device("cpu")],
+        [torch.device("cpu"), torch.device("cuda")],
+        [torch.device("cuda"), torch.device("cpu")],
+        [torch.device("cuda"), torch.device("cuda")],
+    ],
+)
+def test_ssim_device(available_device, metric_device, y_pred_device):
+    if available_device == "cpu":
+        pytest.skip("This test requires a cuda device.")
+
+    data_range = 1.0
+    sigma = 1.5
+    shape = (12, 5, 256, 256)
+
+    ssim = SSIM(data_range=data_range, sigma=sigma, device=metric_device)
+
+    y_pred = torch.rand(shape, device=y_pred_device)
+    y = y_pred * 0.8
+
+    if metric_device == torch.device("cuda") and y_pred_device == torch.device("cpu"):
+        with pytest.warns(UserWarning):
+            ssim.update((y_pred, y))
+    else:
+        ssim.update((y_pred, y))
+
+    if metric_device == torch.device("cuda") or y_pred_device == torch.device("cuda"):
+        # A tensor will always have the device index set
+        excepted_device = torch.device("cuda:0")
+    else:
+        excepted_device = torch.device("cpu")
+
+    assert ssim._kernel.device == excepted_device
+
+
 def test_ssim_variable_batchsize(available_device):
     # Checks https://github.com/pytorch/ignite/issues/2532
     sigma = 1.5
@@ -128,6 +191,21 @@ def test_ssim_variable_batchsize(available_device):
     assert np.allclose(out, expected)
 
 
+def test_ssim_variable_channel(available_device):
+    y_preds = [
+        torch.rand(12, 5, 28, 28, device=available_device),
+        torch.rand(12, 4, 28, 28, device=available_device),
+        torch.rand(12, 7, 28, 28, device=available_device),
+        torch.rand(12, 3, 28, 28, device=available_device),
+        torch.rand(12, 11, 28, 28, device=available_device),
+        torch.rand(12, 6, 28, 28, device=available_device),
+    ]
+    y_true = [v * 0.8 for v in y_preds]
+
+    for y_pred, y in zip(y_preds, y_true):
+        compare_ssim_ignite_skiimg(y_pred, y, available_device)
+
+
 @pytest.mark.parametrize(
     "dtype, precision", [(torch.bfloat16, 2e-3), (torch.float16, 4e-4), (torch.float32, 2e-5), (torch.float64, 2e-5)]
 )
@@ -136,7 +214,12 @@ def test_cuda_ssim_dtypes(available_device, dtype, precision):
     if available_device == "cpu" and dtype in [torch.float16, torch.bfloat16]:
         pytest.skip(reason=f"Unsupported dtype {dtype} on CPU device")
 
-    test_ssim(available_device, (12, 3, 28, 28), 11, True, False, dtype=dtype, precision=precision)
+    shape = (12, 3, 28, 28)
+
+    y_pred = torch.rand(shape, device=available_device, dtype=dtype)
+    y = y_pred * 0.8
+
+    compare_ssim_ignite_skiimg(y_pred, y, available_device, precision)
 
 
 @pytest.mark.parametrize("metric_device", ["cpu", "process_device"])