We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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 ?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `ignite/metrics/ssim.py`
Content:
```
1 from typing import Callable, Sequence, Union
2 
3 import torch
4 import torch.nn.functional as F
5 
6 from ignite.exceptions import NotComputableError
7 from ignite.metrics.metric import Metric, reinit__is_reduced, sync_all_reduce
8 
9 __all__ = ["SSIM"]
10 
11 
12 class SSIM(Metric):
13     """
14     Computes Structual Similarity Index Measure
15 
16     Args:
17         data_range (int or float): Range of the image. Typically, ``1.0`` or ``255``.
18         kernel_size (int or list or tuple of int): Size of the kernel. Default: (11, 11)
19         sigma (float or list or tuple of float): Standard deviation of the gaussian kernel.
20             Argument is used if ``gaussian=True``. Default: (1.5, 1.5)
21         k1 (float): Parameter of SSIM. Default: 0.01
22         k2 (float): Parameter of SSIM. Default: 0.03
23         gaussian (bool): ``True`` to use gaussian kernel, ``False`` to use uniform kernel
24         output_transform (callable, optional): A callable that is used to transform the
25             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
26             form expected by the metric.
27 
28     Example:
29 
30     To use with ``Engine`` and ``process_function``, simply attach the metric instance to the engine.
31     The output of the engine's ``process_function`` needs to be in the format of
32     ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y, ...}``.
33 
34     ``y_pred`` and ``y`` can be un-normalized or normalized image tensors. Depending on that, the user might need
35     to adjust ``data_range``. ``y_pred`` and ``y`` should have the same shape.
36 
37     .. code-block:: python
38 
39         def process_function(engine, batch):
40             # ...
41             return y_pred, y
42         engine = Engine(process_function)
43         metric = SSIM(data_range=1.0)
44         metric.attach(engine, "ssim")
45     """
46 
47     def __init__(
48         self,
49         data_range: Union[int, float],
50         kernel_size: Union[int, Sequence[int]] = (11, 11),
51         sigma: Union[float, Sequence[float]] = (1.5, 1.5),
52         k1: float = 0.01,
53         k2: float = 0.03,
54         gaussian: bool = True,
55         output_transform: Callable = lambda x: x,
56     ):
57         if isinstance(kernel_size, int):
58             self.kernel_size = [kernel_size, kernel_size]
59         elif isinstance(kernel_size, Sequence):
60             self.kernel_size = kernel_size
61         else:
62             raise ValueError("Argument kernel_size should be either int or a sequence of int.")
63 
64         if isinstance(sigma, float):
65             self.sigma = [sigma, sigma]
66         elif isinstance(sigma, Sequence):
67             self.sigma = sigma
68         else:
69             raise ValueError("Argument sigma should be either float or a sequence of float.")
70 
71         if any(x % 2 == 0 or x <= 0 for x in self.kernel_size):
72             raise ValueError("Expected kernel_size to have odd positive number. Got {}.".format(kernel_size))
73 
74         if any(y <= 0 for y in self.sigma):
75             raise ValueError("Expected sigma to have positive number. Got {}.".format(sigma))
76 
77         self.gaussian = gaussian
78         self.c1 = (k1 * data_range) ** 2
79         self.c2 = (k2 * data_range) ** 2
80         self.pad_h = (self.kernel_size[0] - 1) // 2
81         self.pad_w = (self.kernel_size[1] - 1) // 2
82         self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)
83         super(SSIM, self).__init__(output_transform=output_transform)
84 
85     @reinit__is_reduced
86     def reset(self) -> None:
87         self._sum_of_batchwise_ssim = 0.0
88         self._num_examples = 0
89         self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)
90 
91     def _uniform(self, kernel_size):
92         max, min = 2.5, -2.5
93         kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)
94         for i, j in enumerate(kernel):
95             if min <= j <= max:
96                 kernel[i] = 1 / (max - min)
97             else:
98                 kernel[i] = 0
99 
100         return kernel.unsqueeze(dim=0)  # (1, kernel_size)
101 
102     def _gaussian(self, kernel_size, sigma):
103         kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)
104         gauss = torch.exp(-kernel.pow(2) / (2 * pow(sigma, 2)))
105         return (gauss / gauss.sum()).unsqueeze(dim=0)  # (1, kernel_size)
106 
107     def _gaussian_or_uniform_kernel(self, kernel_size, sigma):
108         if self.gaussian:
109             kernel_x = self._gaussian(kernel_size[0], sigma[0])
110             kernel_y = self._gaussian(kernel_size[1], sigma[1])
111         else:
112             kernel_x = self._uniform(kernel_size[0])
113             kernel_y = self._uniform(kernel_size[1])
114 
115         return torch.matmul(kernel_x.t(), kernel_y)  # (kernel_size, 1) * (1, kernel_size)
116 
117     @reinit__is_reduced
118     def update(self, output: Sequence[torch.Tensor]) -> None:
119         y_pred, y = output[0].detach(), output[1].detach()
120 
121         if y_pred.dtype != y.dtype:
122             raise TypeError(
123                 "Expected y_pred and y to have the same data type. Got y_pred: {} and y: {}.".format(
124                     y_pred.dtype, y.dtype
125                 )
126             )
127 
128         if y_pred.shape != y.shape:
129             raise ValueError(
130                 "Expected y_pred and y to have the same shape. Got y_pred: {} and y: {}.".format(y_pred.shape, y.shape)
131             )
132 
133         if len(y_pred.shape) != 4 or len(y.shape) != 4:
134             raise ValueError(
135                 "Expected y_pred and y to have BxCxHxW shape. Got y_pred: {} and y: {}.".format(y_pred.shape, y.shape)
136             )
137 
138         channel = y_pred.size(1)
139         if len(self._kernel.shape) < 4:
140             self._kernel = self._kernel.expand(channel, 1, -1, -1).to(device=y_pred.device)
141 
142         y_pred = F.pad(y_pred, (self.pad_w, self.pad_w, self.pad_h, self.pad_h), mode="reflect")
143         y = F.pad(y, (self.pad_w, self.pad_w, self.pad_h, self.pad_h), mode="reflect")
144 
145         input_list = torch.cat([y_pred, y, y_pred * y_pred, y * y, y_pred * y])
146         outputs = F.conv2d(input_list, self._kernel, groups=channel)
147 
148         output_list = [outputs[x * y_pred.size(0) : (x + 1) * y_pred.size(0)] for x in range(len(outputs))]
149 
150         mu_pred_sq = output_list[0].pow(2)
151         mu_target_sq = output_list[1].pow(2)
152         mu_pred_target = output_list[0] * output_list[1]
153 
154         sigma_pred_sq = output_list[2] - mu_pred_sq
155         sigma_target_sq = output_list[3] - mu_target_sq
156         sigma_pred_target = output_list[4] - mu_pred_target
157 
158         a1 = 2 * mu_pred_target + self.c1
159         a2 = 2 * sigma_pred_target + self.c2
160         b1 = mu_pred_sq + mu_target_sq + self.c1
161         b2 = sigma_pred_sq + sigma_target_sq + self.c2
162 
163         ssim_idx = (a1 * a2) / (b1 * b2)
164         self._sum_of_batchwise_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64)
165         self._num_examples += y.shape[0]
166 
167     @sync_all_reduce("_sum_of_batchwise_ssim", "_num_examples")
168     def compute(self) -> torch.Tensor:
169         if self._num_examples == 0:
170             raise NotComputableError("SSIM must have at least one example before it can be computed.")
171         return torch.sum(self._sum_of_batchwise_ssim / self._num_examples)
172
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/ignite/metrics/ssim.py b/ignite/metrics/ssim.py
--- a/ignite/metrics/ssim.py
+++ b/ignite/metrics/ssim.py
@@ -53,6 +53,7 @@
         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 @@
         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 @@
         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 @@
         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")

{"golden_diff": "diff --git a/ignite/metrics/ssim.py b/ignite/metrics/ssim.py\n--- a/ignite/metrics/ssim.py\n+++ b/ignite/metrics/ssim.py\n@@ -53,6 +53,7 @@\n         k2: float = 0.03,\n         gaussian: bool = True,\n         output_transform: Callable = lambda x: x,\n+        device: Union[str, torch.device] = torch.device(\"cpu\"),\n     ):\n         if isinstance(kernel_size, int):\n             self.kernel_size = [kernel_size, kernel_size]\n@@ -74,23 +75,24 @@\n         if any(y <= 0 for y in self.sigma):\n             raise ValueError(\"Expected sigma to have positive number. Got {}.\".format(sigma))\n \n+        super(SSIM, self).__init__(output_transform=output_transform, device=device)\n         self.gaussian = gaussian\n         self.c1 = (k1 * data_range) ** 2\n         self.c2 = (k2 * data_range) ** 2\n         self.pad_h = (self.kernel_size[0] - 1) // 2\n         self.pad_w = (self.kernel_size[1] - 1) // 2\n         self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)\n-        super(SSIM, self).__init__(output_transform=output_transform)\n \n     @reinit__is_reduced\n     def reset(self) -> None:\n-        self._sum_of_batchwise_ssim = 0.0\n+        self._sum_of_batchwise_ssim = 0.0  # Not a tensor because batch size is not known in advance.\n         self._num_examples = 0\n         self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)\n \n     def _uniform(self, kernel_size):\n         max, min = 2.5, -2.5\n-        kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)\n+        ksize_half = (kernel_size - 1) * 0.5\n+        kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size, device=self._device)\n         for i, j in enumerate(kernel):\n             if min <= j <= max:\n                 kernel[i] = 1 / (max - min)\n@@ -100,8 +102,9 @@\n         return kernel.unsqueeze(dim=0)  # (1, kernel_size)\n \n     def _gaussian(self, kernel_size, sigma):\n-        kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)\n-        gauss = torch.exp(-kernel.pow(2) / (2 * pow(sigma, 2)))\n+        ksize_half = (kernel_size - 1) * 0.5\n+        kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size, device=self._device)\n+        gauss = torch.exp(-0.5 * (kernel / sigma).pow(2))\n         return (gauss / gauss.sum()).unsqueeze(dim=0)  # (1, kernel_size)\n \n     def _gaussian_or_uniform_kernel(self, kernel_size, sigma):\n@@ -161,7 +164,7 @@\n         b2 = sigma_pred_sq + sigma_target_sq + self.c2\n \n         ssim_idx = (a1 * a2) / (b1 * b2)\n-        self._sum_of_batchwise_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64)\n+        self._sum_of_batchwise_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64).to(self._device)\n         self._num_examples += y.shape[0]\n \n     @sync_all_reduce(\"_sum_of_batchwise_ssim\", \"_num_examples\")\n", "issue": "Update SSIM metric\n## \ud83d\ude80 Feature\r\n\r\nCurrent 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.\r\n\r\n@n2cholas , could you please work on that ?\n", "before_files": [{"content": "from typing import Callable, Sequence, Union\n\nimport torch\nimport torch.nn.functional as F\n\nfrom ignite.exceptions import NotComputableError\nfrom ignite.metrics.metric import Metric, reinit__is_reduced, sync_all_reduce\n\n__all__ = [\"SSIM\"]\n\n\nclass SSIM(Metric):\n    \"\"\"\n    Computes Structual Similarity Index Measure\n\n    Args:\n        data_range (int or float): Range of the image. Typically, ``1.0`` or ``255``.\n        kernel_size (int or list or tuple of int): Size of the kernel. Default: (11, 11)\n        sigma (float or list or tuple of float): Standard deviation of the gaussian kernel.\n            Argument is used if ``gaussian=True``. Default: (1.5, 1.5)\n        k1 (float): Parameter of SSIM. Default: 0.01\n        k2 (float): Parameter of SSIM. Default: 0.03\n        gaussian (bool): ``True`` to use gaussian kernel, ``False`` to use uniform kernel\n        output_transform (callable, optional): A callable that is used to transform the\n            :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the\n            form expected by the metric.\n\n    Example:\n\n    To use with ``Engine`` and ``process_function``, simply attach the metric instance to the engine.\n    The output of the engine's ``process_function`` needs to be in the format of\n    ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y, ...}``.\n\n    ``y_pred`` and ``y`` can be un-normalized or normalized image tensors. Depending on that, the user might need\n    to adjust ``data_range``. ``y_pred`` and ``y`` should have the same shape.\n\n    .. code-block:: python\n\n        def process_function(engine, batch):\n            # ...\n            return y_pred, y\n        engine = Engine(process_function)\n        metric = SSIM(data_range=1.0)\n        metric.attach(engine, \"ssim\")\n    \"\"\"\n\n    def __init__(\n        self,\n        data_range: Union[int, float],\n        kernel_size: Union[int, Sequence[int]] = (11, 11),\n        sigma: Union[float, Sequence[float]] = (1.5, 1.5),\n        k1: float = 0.01,\n        k2: float = 0.03,\n        gaussian: bool = True,\n        output_transform: Callable = lambda x: x,\n    ):\n        if isinstance(kernel_size, int):\n            self.kernel_size = [kernel_size, kernel_size]\n        elif isinstance(kernel_size, Sequence):\n            self.kernel_size = kernel_size\n        else:\n            raise ValueError(\"Argument kernel_size should be either int or a sequence of int.\")\n\n        if isinstance(sigma, float):\n            self.sigma = [sigma, sigma]\n        elif isinstance(sigma, Sequence):\n            self.sigma = sigma\n        else:\n            raise ValueError(\"Argument sigma should be either float or a sequence of float.\")\n\n        if any(x % 2 == 0 or x <= 0 for x in self.kernel_size):\n            raise ValueError(\"Expected kernel_size to have odd positive number. Got {}.\".format(kernel_size))\n\n        if any(y <= 0 for y in self.sigma):\n            raise ValueError(\"Expected sigma to have positive number. Got {}.\".format(sigma))\n\n        self.gaussian = gaussian\n        self.c1 = (k1 * data_range) ** 2\n        self.c2 = (k2 * data_range) ** 2\n        self.pad_h = (self.kernel_size[0] - 1) // 2\n        self.pad_w = (self.kernel_size[1] - 1) // 2\n        self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)\n        super(SSIM, self).__init__(output_transform=output_transform)\n\n    @reinit__is_reduced\n    def reset(self) -> None:\n        self._sum_of_batchwise_ssim = 0.0\n        self._num_examples = 0\n        self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)\n\n    def _uniform(self, kernel_size):\n        max, min = 2.5, -2.5\n        kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)\n        for i, j in enumerate(kernel):\n            if min <= j <= max:\n                kernel[i] = 1 / (max - min)\n            else:\n                kernel[i] = 0\n\n        return kernel.unsqueeze(dim=0)  # (1, kernel_size)\n\n    def _gaussian(self, kernel_size, sigma):\n        kernel = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=torch.float32)\n        gauss = torch.exp(-kernel.pow(2) / (2 * pow(sigma, 2)))\n        return (gauss / gauss.sum()).unsqueeze(dim=0)  # (1, kernel_size)\n\n    def _gaussian_or_uniform_kernel(self, kernel_size, sigma):\n        if self.gaussian:\n            kernel_x = self._gaussian(kernel_size[0], sigma[0])\n            kernel_y = self._gaussian(kernel_size[1], sigma[1])\n        else:\n            kernel_x = self._uniform(kernel_size[0])\n            kernel_y = self._uniform(kernel_size[1])\n\n        return torch.matmul(kernel_x.t(), kernel_y)  # (kernel_size, 1) * (1, kernel_size)\n\n    @reinit__is_reduced\n    def update(self, output: Sequence[torch.Tensor]) -> None:\n        y_pred, y = output[0].detach(), output[1].detach()\n\n        if y_pred.dtype != y.dtype:\n            raise TypeError(\n                \"Expected y_pred and y to have the same data type. Got y_pred: {} and y: {}.\".format(\n                    y_pred.dtype, y.dtype\n                )\n            )\n\n        if y_pred.shape != y.shape:\n            raise ValueError(\n                \"Expected y_pred and y to have the same shape. Got y_pred: {} and y: {}.\".format(y_pred.shape, y.shape)\n            )\n\n        if len(y_pred.shape) != 4 or len(y.shape) != 4:\n            raise ValueError(\n                \"Expected y_pred and y to have BxCxHxW shape. Got y_pred: {} and y: {}.\".format(y_pred.shape, y.shape)\n            )\n\n        channel = y_pred.size(1)\n        if len(self._kernel.shape) < 4:\n            self._kernel = self._kernel.expand(channel, 1, -1, -1).to(device=y_pred.device)\n\n        y_pred = F.pad(y_pred, (self.pad_w, self.pad_w, self.pad_h, self.pad_h), mode=\"reflect\")\n        y = F.pad(y, (self.pad_w, self.pad_w, self.pad_h, self.pad_h), mode=\"reflect\")\n\n        input_list = torch.cat([y_pred, y, y_pred * y_pred, y * y, y_pred * y])\n        outputs = F.conv2d(input_list, self._kernel, groups=channel)\n\n        output_list = [outputs[x * y_pred.size(0) : (x + 1) * y_pred.size(0)] for x in range(len(outputs))]\n\n        mu_pred_sq = output_list[0].pow(2)\n        mu_target_sq = output_list[1].pow(2)\n        mu_pred_target = output_list[0] * output_list[1]\n\n        sigma_pred_sq = output_list[2] - mu_pred_sq\n        sigma_target_sq = output_list[3] - mu_target_sq\n        sigma_pred_target = output_list[4] - mu_pred_target\n\n        a1 = 2 * mu_pred_target + self.c1\n        a2 = 2 * sigma_pred_target + self.c2\n        b1 = mu_pred_sq + mu_target_sq + self.c1\n        b2 = sigma_pred_sq + sigma_target_sq + self.c2\n\n        ssim_idx = (a1 * a2) / (b1 * b2)\n        self._sum_of_batchwise_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64)\n        self._num_examples += y.shape[0]\n\n    @sync_all_reduce(\"_sum_of_batchwise_ssim\", \"_num_examples\")\n    def compute(self) -> torch.Tensor:\n        if self._num_examples == 0:\n            raise NotComputableError(\"SSIM must have at least one example before it can be computed.\")\n        return torch.sum(self._sum_of_batchwise_ssim / self._num_examples)\n", "path": "ignite/metrics/ssim.py"}], "after_files": [{"content": "from typing import Callable, Sequence, Union\n\nimport torch\nimport torch.nn.functional as F\n\nfrom ignite.exceptions import NotComputableError\nfrom ignite.metrics.metric import Metric, reinit__is_reduced, sync_all_reduce\n\n__all__ = [\"SSIM\"]\n\n\nclass SSIM(Metric):\n    \"\"\"\n    Computes Structual Similarity Index Measure\n\n    Args:\n        data_range (int or float): Range of the image. Typically, ``1.0`` or ``255``.\n        kernel_size (int or list or tuple of int): Size of the kernel. Default: (11, 11)\n        sigma (float or list or tuple of float): Standard deviation of the gaussian kernel.\n            Argument is used if ``gaussian=True``. Default: (1.5, 1.5)\n        k1 (float): Parameter of SSIM. Default: 0.01\n        k2 (float): Parameter of SSIM. Default: 0.03\n        gaussian (bool): ``True`` to use gaussian kernel, ``False`` to use uniform kernel\n        output_transform (callable, optional): A callable that is used to transform the\n            :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the\n            form expected by the metric.\n\n    Example:\n\n    To use with ``Engine`` and ``process_function``, simply attach the metric instance to the engine.\n    The output of the engine's ``process_function`` needs to be in the format of\n    ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y, ...}``.\n\n    ``y_pred`` and ``y`` can be un-normalized or normalized image tensors. Depending on that, the user might need\n    to adjust ``data_range``. ``y_pred`` and ``y`` should have the same shape.\n\n    .. code-block:: python\n\n        def process_function(engine, batch):\n            # ...\n            return y_pred, y\n        engine = Engine(process_function)\n        metric = SSIM(data_range=1.0)\n        metric.attach(engine, \"ssim\")\n    \"\"\"\n\n    def __init__(\n        self,\n        data_range: Union[int, float],\n        kernel_size: Union[int, Sequence[int]] = (11, 11),\n        sigma: Union[float, Sequence[float]] = (1.5, 1.5),\n        k1: float = 0.01,\n        k2: float = 0.03,\n        gaussian: bool = True,\n        output_transform: Callable = lambda x: x,\n        device: Union[str, torch.device] = torch.device(\"cpu\"),\n    ):\n        if isinstance(kernel_size, int):\n            self.kernel_size = [kernel_size, kernel_size]\n        elif isinstance(kernel_size, Sequence):\n            self.kernel_size = kernel_size\n        else:\n            raise ValueError(\"Argument kernel_size should be either int or a sequence of int.\")\n\n        if isinstance(sigma, float):\n            self.sigma = [sigma, sigma]\n        elif isinstance(sigma, Sequence):\n            self.sigma = sigma\n        else:\n            raise ValueError(\"Argument sigma should be either float or a sequence of float.\")\n\n        if any(x % 2 == 0 or x <= 0 for x in self.kernel_size):\n            raise ValueError(\"Expected kernel_size to have odd positive number. Got {}.\".format(kernel_size))\n\n        if any(y <= 0 for y in self.sigma):\n            raise ValueError(\"Expected sigma to have positive number. Got {}.\".format(sigma))\n\n        super(SSIM, self).__init__(output_transform=output_transform, device=device)\n        self.gaussian = gaussian\n        self.c1 = (k1 * data_range) ** 2\n        self.c2 = (k2 * data_range) ** 2\n        self.pad_h = (self.kernel_size[0] - 1) // 2\n        self.pad_w = (self.kernel_size[1] - 1) // 2\n        self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)\n\n    @reinit__is_reduced\n    def reset(self) -> None:\n        self._sum_of_batchwise_ssim = 0.0  # Not a tensor because batch size is not known in advance.\n        self._num_examples = 0\n        self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma)\n\n    def _uniform(self, kernel_size):\n        max, min = 2.5, -2.5\n        ksize_half = (kernel_size - 1) * 0.5\n        kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size, device=self._device)\n        for i, j in enumerate(kernel):\n            if min <= j <= max:\n                kernel[i] = 1 / (max - min)\n            else:\n                kernel[i] = 0\n\n        return kernel.unsqueeze(dim=0)  # (1, kernel_size)\n\n    def _gaussian(self, kernel_size, sigma):\n        ksize_half = (kernel_size - 1) * 0.5\n        kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size, device=self._device)\n        gauss = torch.exp(-0.5 * (kernel / sigma).pow(2))\n        return (gauss / gauss.sum()).unsqueeze(dim=0)  # (1, kernel_size)\n\n    def _gaussian_or_uniform_kernel(self, kernel_size, sigma):\n        if self.gaussian:\n            kernel_x = self._gaussian(kernel_size[0], sigma[0])\n            kernel_y = self._gaussian(kernel_size[1], sigma[1])\n        else:\n            kernel_x = self._uniform(kernel_size[0])\n            kernel_y = self._uniform(kernel_size[1])\n\n        return torch.matmul(kernel_x.t(), kernel_y)  # (kernel_size, 1) * (1, kernel_size)\n\n    @reinit__is_reduced\n    def update(self, output: Sequence[torch.Tensor]) -> None:\n        y_pred, y = output[0].detach(), output[1].detach()\n\n        if y_pred.dtype != y.dtype:\n            raise TypeError(\n                \"Expected y_pred and y to have the same data type. Got y_pred: {} and y: {}.\".format(\n                    y_pred.dtype, y.dtype\n                )\n            )\n\n        if y_pred.shape != y.shape:\n            raise ValueError(\n                \"Expected y_pred and y to have the same shape. Got y_pred: {} and y: {}.\".format(y_pred.shape, y.shape)\n            )\n\n        if len(y_pred.shape) != 4 or len(y.shape) != 4:\n            raise ValueError(\n                \"Expected y_pred and y to have BxCxHxW shape. Got y_pred: {} and y: {}.\".format(y_pred.shape, y.shape)\n            )\n\n        channel = y_pred.size(1)\n        if len(self._kernel.shape) < 4:\n            self._kernel = self._kernel.expand(channel, 1, -1, -1).to(device=y_pred.device)\n\n        y_pred = F.pad(y_pred, (self.pad_w, self.pad_w, self.pad_h, self.pad_h), mode=\"reflect\")\n        y = F.pad(y, (self.pad_w, self.pad_w, self.pad_h, self.pad_h), mode=\"reflect\")\n\n        input_list = torch.cat([y_pred, y, y_pred * y_pred, y * y, y_pred * y])\n        outputs = F.conv2d(input_list, self._kernel, groups=channel)\n\n        output_list = [outputs[x * y_pred.size(0) : (x + 1) * y_pred.size(0)] for x in range(len(outputs))]\n\n        mu_pred_sq = output_list[0].pow(2)\n        mu_target_sq = output_list[1].pow(2)\n        mu_pred_target = output_list[0] * output_list[1]\n\n        sigma_pred_sq = output_list[2] - mu_pred_sq\n        sigma_target_sq = output_list[3] - mu_target_sq\n        sigma_pred_target = output_list[4] - mu_pred_target\n\n        a1 = 2 * mu_pred_target + self.c1\n        a2 = 2 * sigma_pred_target + self.c2\n        b1 = mu_pred_sq + mu_target_sq + self.c1\n        b2 = sigma_pred_sq + sigma_target_sq + self.c2\n\n        ssim_idx = (a1 * a2) / (b1 * b2)\n        self._sum_of_batchwise_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=torch.float64).to(self._device)\n        self._num_examples += y.shape[0]\n\n    @sync_all_reduce(\"_sum_of_batchwise_ssim\", \"_num_examples\")\n    def compute(self) -> torch.Tensor:\n        if self._num_examples == 0:\n            raise NotComputableError(\"SSIM must have at least one example before it can be computed.\")\n        return torch.sum(self._sum_of_batchwise_ssim / self._num_examples)\n", "path": "ignite/metrics/ssim.py"}]}