Repository: pytorch/ignite (Python)

Issue:
Update SSIM metric
## 🚀 Feature

Current implementation of [SSIM](https://github.com/pytorch/ignite/blob/master/ignite/metrics/ssim.py) does not use torch tensor as internal storage and hasn't device in the constructor. Following #1238 we expect all metrics to have a similar structure.

@n2cholas , could you please work on that ?


Discussion:
@vfdev-5 Will do!



Patch:
diff --git a/ignite/metrics/ssim.py b/ignite/metrics/ssim.py
index 7c20a3260cf7..256ce516b589 100644
--- a/ignite/metrics/ssim.py
+++ b/ignite/metrics/ssim.py
@@ -53,6 +53,7 @@ def __init__(
         k2: float = 0.03,
         gaussian: bool = True,
         output_transform: Callable = lambda x: x,
+        device: Union[str, torch.device] = torch.device("cpu"),
     ):
         if isinstance(kernel_size, int):
             self.kernel_size = [kernel_size, kernel_size]
@@ -74,23 +75,24 @@ def __init__(
         if any(y <= 0 for y in self.sigma):
             raise ValueError("Expected sigma to have positive number. Got {}.".format(sigma))
 
+        super(SSIM, self).__init__(output_transform=output_transform, device=device)
         self.gaussian = gaussian
         self.c1 = (k1 * data_range) ** 2
         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)
-        super(SSIM, self).__init__(output_transform=output_transform)
 
     @reinit__is_reduced
     def reset(self) -> None:
-        self._sum_of_batchwise_ssim = 0.0
+        self._sum_of_batchwise_ssim = 0.0  # Not a tensor because batch size is not known in advance.
         self._num_examples = 0
         self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)
 
     def _uniform(self, kernel_size):
         max, min = 2.5, -2.5
-        kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)
+        ksize_half = (kernel_size - 1) * 0.5
+        kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size, device=self._device)
         for i, j in enumerate(kernel):
             if min <= j <= max:
                 kernel[i] = 1 / (max - min)
@@ -100,8 +102,9 @@ def _uniform(self, kernel_size):
         return kernel.unsqueeze(dim=0)  # (1, kernel_size)
 
     def _gaussian(self, kernel_size, sigma):
-        kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)
-        gauss = torch.exp(-kernel.pow(2) / (2 * pow(sigma, 2)))
+        ksize_half = (kernel_size - 1) * 0.5
+        kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size, device=self._device)
+        gauss = torch.exp(-0.5 * (kernel / sigma).pow(2))
         return (gauss / gauss.sum()).unsqueeze(dim=0)  # (1, kernel_size)
 
     def _gaussian_or_uniform_kernel(self, kernel_size, sigma):
@@ -161,7 +164,7 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
         b2 = sigma_pred_sq + sigma_target_sq + self.c2
 
         ssim_idx = (a1 * a2) / (b1 * b2)
-        self._sum_of_batchwise_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64)
+        self._sum_of_batchwise_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64).to(self._device)
         self._num_examples += y.shape[0]
 
     @sync_all_reduce("_sum_of_batchwise_ssim", "_num_examples")


Test patch:
diff --git a/tests/ignite/metrics/test_ssim.py b/tests/ignite/metrics/test_ssim.py
index 4e69396f8316..b158c181e79b 100644
--- a/tests/ignite/metrics/test_ssim.py
+++ b/tests/ignite/metrics/test_ssim.py
@@ -49,8 +49,8 @@ def test_invalid_ssim():
 
 
 def test_ssim():
-    ssim = SSIM(data_range=1.0)
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    ssim = SSIM(data_range=1.0, device=device)
     y_pred = torch.rand(16, 3, 64, 64, device=device)
     y = y_pred * 0.65
     ssim.update((y_pred, y))
@@ -62,8 +62,8 @@ def test_ssim():
     assert isinstance(ssim.compute(), torch.Tensor)
     assert torch.allclose(ssim.compute(), torch.tensor(np_ssim, dtype=torch.float64, device=device), atol=1e-4)
 
-    ssim = SSIM(data_range=1.0, gaussian=False, kernel_size=7)
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    ssim = SSIM(data_range=1.0, gaussian=False, kernel_size=7, device=device)
     y_pred = torch.rand(16, 3, 227, 227, device=device)
     y = y_pred * 0.65
     ssim.update((y_pred, y))
@@ -84,44 +84,74 @@ def _test_distrib_integration(device, tol=1e-4):
     s = 10
     offset = n_iters * s
 
-    y_pred = torch.rand(offset * idist.get_world_size(), 3, 28, 28, dtype=torch.float, device=device)
-    y = y_pred * 0.65
+    def _test(metric_device):
+        y_pred = torch.rand(offset * idist.get_world_size(), 3, 28, 28, dtype=torch.float, device=device)
+        y = y_pred * 0.65
 
-    def update(engine, i):
-        return (
-            y_pred[i * s + offset * rank : (i + 1) * s + offset * rank],
-            y[i * s + offset * rank : (i + 1) * s + offset * rank],
-        )
+        def update(engine, i):
+            return (
+                y_pred[i * s + offset * rank : (i + 1) * s + offset * rank],
+                y[i * s + offset * rank : (i + 1) * s + offset * rank],
+            )
 
-    engine = Engine(update)
-    SSIM(data_range=1.0).attach(engine, "ssim")
+        engine = Engine(update)
+        SSIM(data_range=1.0, device=metric_device).attach(engine, "ssim")
 
-    data = list(range(n_iters))
-    engine.run(data=data, max_epochs=1)
+        data = list(range(n_iters))
+        engine.run(data=data, max_epochs=1)
 
-    assert "ssim" in engine.state.metrics
-    res = engine.state.metrics["ssim"]
+        assert "ssim" in engine.state.metrics
+        res = engine.state.metrics["ssim"]
 
-    np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy()
-    np_true = np_pred * 0.65
-    true_res = ski_ssim(np_pred, np_true, win_size=11, multichannel=True, gaussian_weights=True, data_range=1.0)
+        np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy()
+        np_true = np_pred * 0.65
+        true_res = ski_ssim(np_pred, np_true, win_size=11, multichannel=True, gaussian_weights=True, data_range=1.0)
 
-    assert pytest.approx(res, abs=tol) == true_res
+        assert pytest.approx(res, abs=tol) == true_res
 
-    engine = Engine(update)
-    SSIM(data_range=1.0, gaussian=False, kernel_size=7).attach(engine, "ssim")
+        engine = Engine(update)
+        SSIM(data_range=1.0, gaussian=False, kernel_size=7, device=metric_device).attach(engine, "ssim")
 
-    data = list(range(n_iters))
-    engine.run(data=data, max_epochs=1)
+        data = list(range(n_iters))
+        engine.run(data=data, max_epochs=1)
 
-    assert "ssim" in engine.state.metrics
-    res = engine.state.metrics["ssim"]
+        assert "ssim" in engine.state.metrics
+        res = engine.state.metrics["ssim"]
+
+        np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy()
+        np_true = np_pred * 0.65
+        true_res = ski_ssim(np_pred, np_true, win_size=7, multichannel=True, gaussian_weights=False, data_range=1.0)
+
+        assert pytest.approx(res, abs=tol) == true_res
+
+    _test("cpu")
+    if torch.device(device).type != "xla":
+        _test(idist.device())
 
-    np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy()
-    np_true = np_pred * 0.65
-    true_res = ski_ssim(np_pred, np_true, win_size=7, multichannel=True, gaussian_weights=False, data_range=1.0)
 
-    assert pytest.approx(res, abs=tol) == true_res
+def _test_distrib_accumulator_device(device):
+
+    metric_devices = [torch.device("cpu")]
+    if torch.device(device).type != "xla":
+        metric_devices.append(idist.device())
+    for metric_device in metric_devices:
+
+        ssim = SSIM(data_range=1.0, device=metric_device)
+        assert ssim._device == metric_device
+        assert ssim._kernel.device == metric_device, "{}:{} vs {}:{}".format(
+            type(ssim._kernel.device), ssim._kernel.device, type(metric_device), metric_device
+        )
+
+        y_pred = torch.rand(4, 3, 28, 28, dtype=torch.float, device=device)
+        y = y_pred * 0.65
+        ssim.update((y_pred, y))
+
+        assert ssim._sum_of_batchwise_ssim.device == metric_device, "{}:{} vs {}:{}".format(
+            type(ssim._sum_of_batchwise_ssim.device),
+            ssim._sum_of_batchwise_ssim.device,
+            type(metric_device),
+            metric_device,
+        )
 
 
 @pytest.mark.distributed
@@ -131,6 +161,7 @@ def test_distrib_gpu(local_rank, distributed_context_single_node_nccl):
 
     device = "cuda:{}".format(local_rank)
     _test_distrib_integration(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.distributed
@@ -138,6 +169,7 @@ def test_distrib_gpu(local_rank, distributed_context_single_node_nccl):
 def test_distrib_cpu(distributed_context_single_node_gloo):
     device = "cpu"
     _test_distrib_integration(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.multinode_distributed
@@ -146,6 +178,7 @@ def test_distrib_cpu(distributed_context_single_node_gloo):
 def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
     device = "cpu"
     _test_distrib_integration(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.multinode_distributed
@@ -154,6 +187,7 @@ def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
 def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
     device = "cuda:{}".format(distributed_context_multi_node_nccl["local_rank"])
     _test_distrib_integration(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu
@@ -162,11 +196,13 @@ def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
 def test_distrib_single_device_xla():
     device = idist.device()
     _test_distrib_integration(device, tol=1e-3)
+    _test_distrib_accumulator_device(device)
 
 
 def _test_distrib_xla_nprocs(index):
     device = idist.device()
     _test_distrib_integration(device, tol=1e-3)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu
@@ -186,3 +222,4 @@ def test_distrib_hvd(gloo_hvd_executor):
     nproc = 4 if not torch.cuda.is_available() else torch.cuda.device_count()
 
     gloo_hvd_executor(_test_distrib_integration, (device,), np=nproc, do_init=True)
+    gloo_hvd_executor(_test_distrib_accumulator_device, (device,), np=nproc, do_init=True)