🩹 Fix performance issues due to torch.linalg.norm

Author: francois-rozet

piqa/fsim.py

@@ -28,6 +28,7 @@
     filter_grid,
     log_gabor,
     channel_conv,
+    l2_norm,
 )
 
 
@@ -87,8 +88,8 @@ def fsim(
     # Gradient magnitude similarity
     pad = kernel.size(-1) // 2
 
-    g_x = torch.linalg.norm(channel_conv(y_x, kernel, padding=pad), dim=1)
-    g_y = torch.linalg.norm(channel_conv(y_y, kernel, padding=pad), dim=1)
+    g_x = l2_norm(channel_conv(y_x, kernel, padding=pad), dims=[1])
+    g_y = l2_norm(channel_conv(y_y, kernel, padding=pad), dims=[1])
 
     s_g = (2 * g_x * g_y + t2) / (g_x ** 2 + g_y ** 2 + t2)
 

piqa/gmsd.py

@@ -24,6 +24,7 @@
     prewitt_kernel,
     gradient_kernel,
     channel_conv,
+    l2_norm,
 )
 
 
@@ -75,8 +76,8 @@ def gmsd(
     # Gradient magnitude
     pad = kernel.size(-1) // 2
 
-    gm_x = torch.linalg.norm(channel_conv(x, kernel, padding=pad), dim=1)
-    gm_y = torch.linalg.norm(channel_conv(y, kernel, padding=pad), dim=1)
+    gm_x = l2_norm(channel_conv(x, kernel, padding=pad), dims=[1])
+    gm_y = l2_norm(channel_conv(y, kernel, padding=pad), dims=[1])
 
     gm_xy = gm_x * gm_y
 
@@ -109,7 +110,8 @@ def ms_gmsd(
     without color space conversion.
 
     .. math::
-        \text{MS-GMSD}(x, y) = \sum^{M}_{i = 1} w_i \text{GMSD}(x^i, y^i)
+        \text{MS-GMSD}(x, y) =
+            \sqrt{\sum^{M}_{i = 1} w_i \text{GMSD}(x^i, y^i) ** 2}
 
     where :math:`x^i` and :math:`y^i` are obtained by downsampling
     the initial tensors by a factor :math:`2^{i - 1}`.
@@ -150,8 +152,8 @@ def ms_gmsd(
             c=c, alpha=alpha,
         ))
 
-    msgmsd = torch.stack(gmsds, dim=-1) ** 2
-    msgmsd = torch.sqrt((msgmsd * weights).sum(dim=-1))
+    msgmsd = weights * torch.stack(gmsds, dim=-1) ** 2
+    msgmsd = msgmsd.sum(dim=-1).sqrt()
 
     return msgmsd
 

piqa/lpips.py

@@ -22,6 +22,7 @@
 from typing import Dict, List
 
 from .utils import _jit, assert_type, reduce_tensor
+from .utils.functional import l2_norm
 
 
 ORIGIN: str = 'https://github.com/richzhang/PerceptualSimilarity'
@@ -209,8 +210,8 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
         residuals = []
 
         for lin, fx, fy in zip(self.lins, self.net(input), self.net(target)):
-            fx = fx / torch.linalg.norm(fx, dim=1, keepdim=True)
-            fy = fy / torch.linalg.norm(fy, dim=1, keepdim=True)
+            fx = fx / l2_norm(fx, dims=[1], keepdim=True)
+            fy = fy / l2_norm(fy, dims=[1], keepdim=True)
 
             mse = ((fx - fy) ** 2).mean(dim=(-1, -2), keepdim=True)
             residuals.append(lin(mse).flatten())

piqa/mdsi.py

@@ -22,6 +22,7 @@
     prewitt_kernel,
     gradient_kernel,
     channel_conv,
+    l2_norm,
 )
 
 
@@ -78,12 +79,9 @@ def mdsi(
     # Gradient magnitude
     pad = kernel.size(-1) // 2
 
-    gm_x = torch.linalg.norm(channel_conv(l_x, kernel, padding=pad), dim=1)
-    gm_y = torch.linalg.norm(channel_conv(l_y, kernel, padding=pad), dim=1)
-    gm_avg = torch.linalg.norm(
-        channel_conv((l_x + l_y) / 2., kernel, padding=pad),
-        dim=1,
-    )
+    gm_x = l2_norm(channel_conv(l_x, kernel, padding=pad), dims=[1])
+    gm_y = l2_norm(channel_conv(l_y, kernel, padding=pad), dims=[1])
+    gm_avg = l2_norm(channel_conv(l_x + l_y, kernel, padding=pad), dims=[1]) / 2
 
     gm_x_sq, gm_y_sq, gm_avg_sq = gm_x ** 2, gm_y ** 2, gm_avg ** 2
 

piqa/ssim.py

@@ -187,10 +187,10 @@ def ms_ssim(
 
         css.append(torch.relu(cs) if i + 1 < m else torch.relu(ss))
 
-    msss = torch.stack(css, dim=-1)
-    msss = (msss ** weights).prod(dim=-1)
+    msss = torch.stack(css, dim=-1) ** weights
+    msss = msss.prod(dim=-1).mean(dim=-1)
 
-    return msss.mean(dim=-1)
+    return msss
 
 
 class SSIM(nn.Module):

piqa/utils/functional.py

@@ -335,3 +335,38 @@ def log_gabor(f: Tensor, f_0: float, sigma_f: float) -> Tensor:
     """
 
     return torch.exp(- (f / f_0).log() ** 2 / (2 * sigma_f ** 2))
+
+
+def l2_norm(
+    x: torch.Tensor,
+    dims: List[int],
+    keepdim: bool = False,
+) -> torch.Tensor:
+    r"""Returns the :math:`L_2` norm of :math:`x`.
+
+    .. math:
+        L_2(x) = \left\| x \right\|_2 = \sqrt{\sum_i x^2_i}
+
+    Args:
+        x: A tensor, :math:`(*,)`.
+        dims: The dimensions along which to calculate the norm.
+        keepdim: Whether the output tensor has `dims` retained or not.
+
+    Wikipedia:
+        https://en.wikipedia.org/wiki/Norm_(mathematics)
+
+    Example:
+        >>> x = torch.arange(9).float().view(3, 3)
+        >>> x
+        tensor([[0., 1., 2.],
+                [3., 4., 5.],
+                [6., 7., 8.]])
+        >>> l2_norm(x, dims=[0])
+        tensor([6.7082, 8.1240, 9.6437])
+    """
+
+    x = x ** 2
+    x = x.sum(dim=dims, keepdim=keepdim)
+    x = x.sqrt()
+
+    return x

piqa/vsi.py

@@ -30,6 +30,7 @@
     filter_grid,
     log_gabor,
     channel_conv,
+    l2_norm,
 )
 
 
@@ -87,8 +88,8 @@ def vsi(
     # Gradient magnitude similarity
     pad = kernel.size(-1) // 2
 
-    g_x = torch.linalg.norm(channel_conv(l_x, kernel, padding=pad), dim=1)
-    g_y = torch.linalg.norm(channel_conv(l_y, kernel, padding=pad), dim=1)
+    g_x = l2_norm(channel_conv(l_x, kernel, padding=pad), dims=[1])
+    g_y = l2_norm(channel_conv(l_y, kernel, padding=pad), dims=[1])
 
     s_g = (2 * g_x * g_y + c2) / (g_x ** 2 + g_y ** 2 + c2)
 
@@ -171,7 +172,7 @@ def sdsp(
     x_f = fft.ifft2(fft.fft2(x_lab) * filtr)
     x_f = cx.real(torch.view_as_real(x_f))
 
-    s_f = torch.linalg.norm(x_f, dim=1)
+    s_f = l2_norm(x_f, dims=[1])
 
     # Color prior
     x_ab = x_lab[:, 1:]

setup.py

@@ -10,7 +10,7 @@
 
 setuptools.setup(
     name='piqa',
-    version='1.2.0',
+    version='1.2.1',
     packages=setuptools.find_packages(),
     description='PyTorch Image Quality Assessment',
     keywords='image quality processing metrics torch vision',