This repository contains an efficient fully-fused implementation of SSIM which is differentiable in nature. There are several factors that contribute to an efficient implementation:

• Convolutions in SSIM are spatially localized leading to fully-fused implementation without touching global memory for intermediate steps.
• Backpropagation through Gaussian Convolution is simply another Gaussian Convolution itself.
• Gaussian Convolutions are separable leading to reduced computation.
• Gaussians are symmetric in nature leading to fewer computations.
• Single convolution pass for multiple statistics.

As per the original SSIM paper, this implementation uses 11x11 sized convolution kernel. The weights for it have been hardcoded and this is another reason for it's speed. This implementation currently only supports 2D images but with variable number of channels and batch size.

Thanks to the contributors, this implementation supports the following GPU architectures:

• NVIDIA GPUs (CUDA).
• AMD GPUs (ROCm).
• Apple Silicon (Metal Performance Shaders).
• Intel GPUs (SYCL).

This project has been tested with:

• PyTorch 2.3.1+cu118 and CUDA 11.8 on Ubuntu 24.04 LTS
• PyTorch 2.4.1+cu124 and CUDA 12.4 on Ubuntu 24.04 LTS
• PyTorch 2.5.1+cu124 and CUDA 12.6 on Windows 11

• PyTorch 2.5.1 on macOS 15.7.1

You must have PyTorch installed with the appropriate backend for your GPU before installing fused-ssim. The installation process requires the backend compilers to be available.

Choose the installation method based on your GPU:

First, ensure you have CUDA Toolkit installed on your system (version 11.8 or 12.x recommended).

# For CUDA 12.4
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124

Verify NVCC (CUDA compiler) is available:

nvcc --version

First, ensure you have ROCm installed on your system (version 5.7 or newer recommended).

# For ROCm 6.1
pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.1

Verify HIP compiler is available:

hipcc --version

Install PyTorch 2.5.1 with MPS backend.

pip install torch torchvision

First, ensure you have Intel oneAPI Base Toolkit installed with DPC++/SYCL compiler support.

# Install PyTorch for Intel XPU
pip install torch torchvision --index-url https://download.pytorch.org/whl/xpu

Verify Intel SYCL compiler is available:

icpx --version

source /opt/intel/oneapi/setvars.sh

git clone https://github.com/rahul-goel/fused-ssim.git
cd fused-ssim
pip install --no-build-isolation .

python -m build --no-isolation --wheel

cmd /k "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
powershell
cmd /k "C:\Program Files (x86)\Intel\oneAPI\setvars.bat"
powershell

Once PyTorch and the appropriate backend compiler are installed:

# Install from GitHub (recommended)
pip install git+https://github.com/rahul-goel/fused-ssim/ --no-build-isolation

# Or clone and install locally
git clone https://github.com/rahul-goel/fused-ssim.git
cd fused-ssim
pip install . --no-build-isolation

The setup.py script will automatically detect your GPU architecture. For verbose output:

pip install git+https://github.com/rahul-goel/fused-ssim/ -v --no-build-isolation

If the above commands don't work, try:

python setup.py install

If you want to specify the GPU architecture manually, like for example to compile a docker image that will run in a different host, you can do so by setting the CUDA_ARCHITECTURES environment variable. For example, to set it to 8.9 and 12.0, run CUDA_ARCHITECTURES="89;120" pip install git+https://github.com/rahul-goel/fused-ssim/.

• CUDA errors: Ensure your CUDA Toolkit version matches your PyTorch CUDA version
• ROCm errors: Verify ROCm installation with rocm-smi and check PyTorch ROCm compatibility
• Metal errors: Ensure Xcode Command Line Tools are installed and up to date
• Intel errors: Source the Intel oneAPI environment before installation: source /opt/intel/oneapi/setvars.sh

• You must have CUDA and PyTorch+CUDA installed in you Python 3.X environment. This project has currently been tested with:
  • PyTorch 2.3.1+cu118 and CUDA 11.8 on Ubuntu 24.04 LTS.
  • PyTorch 2.4.1+cu124 and CUDA 12.4 on Ubuntu 24.04 LTS.
  • PyTorch 2.5.1+cu124 and CUDA 12.6 on Windows 11.
• Run pip install git+https://github.com/rahul-goel/fused-ssim/ --no-build-isolation or clone the repository and run pip install . --no-build-isolation from the root of this project.
• setup.py should detect your GPU architecture automatically. If you want to see the output, run pip install git+https://github.com/rahul-goel/fused-ssim/ -v --no-build-isolation or clone the repository and run pip install . -v --no-build-isolation from the root of this project.
• If you want to specify the GPU architecture manually, like for example to compile a docker image that will run in a different host, you can do so by setting the CUDA_ARCHITECTURES environment variable. For example, to set it to 8.9 and 12.0, run CUDA_ARCHITECTURES="89;120" pip install git+https://github.com/rahul-goel/fused-ssim/.
• If the previous command does not work, run python setup.py install from the root of this project.

import torch
from fused_ssim import fused_ssim

# predicted_image, gt_image: [BS, CH, H, W]
# predicted_image is differentiable
gt_image = torch.rand(2, 3, 1080, 1920)
predicted_image = torch.nn.Parameter(torch.rand_like(gt_image))
ssim_value = fused_ssim(predicted_image, gt_image)

By default, same padding is used. To use valid padding which is the kind of padding used by pytorch-mssim:

ssim_value = fused_ssim(predicted_image, gt_image, padding="valid")

If you don't want to train and use this only for inference, use the following for even faster speed:

with torch.no_grad():
  ssim_value = fused_ssim(predicted_image, gt_image, train=False)

• Currently, only one of the images is allowed to be differentiable i.e. only the first image can be nn.Parameter.
• Limited to 2D images.
• Images must be normalized to range [0, 1].
• Standard 11x11 convolutions supported.

This implementation is 5-8x faster than the previous fastest (to the best of my knowledge) differentiable SSIM implementation pytorch-msssim.

If you leverage fused SSIM for your research work, please cite our main paper:

@inproceedings{taming3dgs,
    author = {Mallick, Saswat Subhajyoti and Goel, Rahul and Kerbl, Bernhard and Steinberger, Markus and Carrasco, Francisco Vicente and De La Torre, Fernando},
    title = {Taming 3DGS: High-Quality Radiance Fields with Limited Resources},
    year = {2024},
    url = {https://doi.org/10.1145/3680528.3687694},
    doi = {10.1145/3680528.3687694},
    booktitle = {SIGGRAPH Asia 2024 Conference Papers},
    series = {SA '24}
}

Thanks to:

• Bernhard for the idea.
• asrathore-ai for adding SYCL kernels.
• Anton Smirnov for adding AMD GPU enablement.
• Jonah J. Newton for Apple MPS kernels.
• Janusch for further CUDA optimizations.
• Florian and Ishaan for testing.