Repository: kornia/kornia

Issue:
Faster implementation of MS-SSIM
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
Current implementation impacts training speed significantly (~40% slowdown compared to pure L1)

## Motivation

<!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
MS-SSIM gives some awesome results, and it'd be nice to have an option to switch to a 
* lower accuracy and faster implementation such as https://link.springer.com/article/10.1007/s11554-010-0170-9
* more optimized implementation

## Pitch

<!-- A clear and concise description of what you want to happen. -->
Faster iteration is always great. Difference between workflow for a task which starts in morning and ends during lunch compared to one that ends in evening is huge.

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
Reduce the scales in MS-SSIM

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->

______________________________________________________________________

#### Consider also to contribute to Kornia universe projects :)

<sub>

- [**Tutorials**](https://github.com/kornia/tutorials): our repository containing the tutorials.

</sub>



Discussion:
@kunaltyagi do you have any suggestion for improvements ? I don't think your statement is fair to compare against L1 since the ssim/msssim is a more complex operator than element-wise subtraction and absolute value. But I agree, let's iterate to optimize if possible.

The way is implemented is via 2d convolutions and the reason for being slow might be because of large kernels (computed from sigmas). The default values gives the last a sigma of  8.0 which makes the max a kernel size of 32 times the number of scales (by default 5) times channels if i understand well the implementation. Definitely a large kernel.

https://github.com/kornia/kornia/blob/master/kornia/losses/ms_ssim.py
I have a couple of ideas, but I don't know how much they are worth since my experience in PyTorch doesn't include bench-marking.
1. Use integral values. They are usually fast than float (though does it matter in pytorch?)
2. Reduce kernel size -> Will have an immediate impact
3. Use jit or compile the kernel to allow optimizations based on the known values of kernel. Eg: multiplication with 5 can be implemented as a shifts + add but multiplication with unknown integer would be more expensive (generic muladd operation). Though some HW actually perform muladd faster than multiplication and shift (ask Intel)
4. Use other operation instead of expensive variance calculation 

Some of these are the same ideas as shown in the paper linked above. Apologies, but I can't find a PDF which I can distribute.
@kunaltyagi 
1) pytorch needs floating point values as requirements to differentiate.
2) We can explore and give a different default, however, as an api it should be customizable to enable more wide research. 
3) Not clear shifts implements the backward properly.
4) I believe variance is part of the formulation of ssim.

Tagging @t-vi  to contribute with ideas as for pytorch optimization here.

As for benchmark, if you want to try some of those ideas i suggest to read this tutorial
https://pytorch.org/tutorials/recipes/recipes/benchmark.html

Eventually, use this template: https://github.com/kornia/kornia/blob/master/test/benchmark.py

It is indeed pretty slow with large kernels which is expected. Current implementation uses regular convolution, but gaussian filter is separable, so using get_gaussian_kernel1d and filter2d_separable gives 10x speedup with the kernel size of 127.

Resolving patch:
diff --git a/kornia/metrics/ssim.py b/kornia/metrics/ssim.py
--- a/kornia/metrics/ssim.py
+++ b/kornia/metrics/ssim.py
@@ -3,7 +3,7 @@
 import torch
 import torch.nn as nn
 
-from kornia.filters import filter2d, get_gaussian_kernel2d
+from kornia.filters import filter2d_separable, get_gaussian_kernel1d
 from kornia.filters.filter import _compute_padding
 
 
@@ -75,19 +75,19 @@ def ssim(
         raise ValueError(f"img1 and img2 shapes must be the same. Got: {img1.shape} and {img2.shape}")
 
     # prepare kernel
-    kernel: torch.Tensor = get_gaussian_kernel2d((window_size, window_size), (1.5, 1.5))
+    kernel: torch.Tensor = get_gaussian_kernel1d(window_size, 1.5, device=img1.device, dtype=img1.dtype)
 
     # compute coefficients
     C1: float = (0.01 * max_val) ** 2
     C2: float = (0.03 * max_val) ** 2
 
     # compute local mean per channel
-    mu1: torch.Tensor = filter2d(img1, kernel)
-    mu2: torch.Tensor = filter2d(img2, kernel)
+    mu1: torch.Tensor = filter2d_separable(img1, kernel, kernel)
+    mu2: torch.Tensor = filter2d_separable(img2, kernel, kernel)
 
     cropping_shape: List[int] = []
     if padding == 'valid':
-        height, width = kernel.shape[-2:]
+        height = width = kernel.shape[-1]
         cropping_shape = _compute_padding([height, width])
         mu1 = _crop(mu1, cropping_shape)
         mu2 = _crop(mu2, cropping_shape)
@@ -98,9 +98,9 @@ def ssim(
     mu2_sq = mu2**2
     mu1_mu2 = mu1 * mu2
 
-    mu_img1_sq = filter2d(img1**2, kernel)
-    mu_img2_sq = filter2d(img2**2, kernel)
-    mu_img1_img2 = filter2d(img1 * img2, kernel)
+    mu_img1_sq = filter2d_separable(img1**2, kernel, kernel)
+    mu_img2_sq = filter2d_separable(img2**2, kernel, kernel)
+    mu_img1_img2 = filter2d_separable(img1 * img2, kernel, kernel)
 
     if padding == 'valid':
         mu_img1_sq = _crop(mu_img1_sq, cropping_shape)


Test patch: