ms_ssim_pytorch

The code was modified from https://github.com/VainF/pytorch-msssim.
Part of the code has been modified to make it faster, takes up less VRAM, and is compatible with pytorch jit.

The dynamic channel version can found here https://github.com/One-sixth/ms_ssim_pytorch/tree/dynamic_channel_num.
More convenient to use but has a little performance loss.

Thanks vegetable09 for finding and fixing a bug that causes gradient nan when ms_ssim backward. #3

If you are using pytorch 1.2, please be careful not to create and destroy this jit module in the training loop (other jit modules may also have this situation), there may be memory leaks. I have tested that pytorch 1.6 does not have this problem. #4

I study to the ssim.py of the library piqa, which makes my implementation of ssim and ms-ssim a little faster than before.

Speed up. Only test on GPU.

losser1 is https://github.com/lizhengwei1992/MS_SSIM_pytorch/blob/master/loss.py 268fc76
losser2 is https://github.com/Po-Hsun-Su/pytorch-ssim/blob/master/pytorch_ssim/__init__.py 881d210
losser3 is https://github.com/VainF/pytorch-msssim/blob/master/pytorch_msssim/ssim.py 5caf547
losser4 is https://github.com/One-sixth/ms_ssim_pytorch/blob/master/ssim.py 1c2f14a
losser5 is https://github.com/francois-rozet/piqa/blob/master/piqa/ssim.py abaf398
https://github.com/francois-rozet/piqa/blob/master/piqa/utils.py abaf398

In pytorch 1.7.1
My test environment: i7-8750H GTX1070-8G

In pytorch 1.1 1.2
My test environment: i7-6700HQ GTX970M-3G

SSIM

Test output

pytorch 1.7.1

Performance Testing SSIM

testing losser2
cuda time 39963.15625
perf_counter time 35.9110169

testing losser3
cuda time 17141.841796875
perf_counter time 17.124456199999997

testing losser4
cuda time 13205.0322265625
perf_counter time 10.477991699999997

testing losser5
cuda time 13142.8232421875
perf_counter time 11.079514100000011

pytorch 1.2

Performance Testing SSIM

testing losser2
cuda time 89290.7734375
perf_counter time 87.1042247

testing losser3
cuda time 36153.64453125
perf_counter time 36.09167939999999

testing losser4
cuda time 31085.455078125
perf_counter time 29.80807200000001

pytorch 1.1

Performance Testing SSIM

testing losser2
cuda time 88990.0703125
perf_counter time 86.80163019999999

testing losser3
cuda time 36119.06640625
perf_counter time 36.057978399999996

testing losser4
cuda time 34708.8359375
perf_counter time 33.916086199999995

MS-SSIM

Test output

pytorch 1.7.1

Performance Testing MS_SSIM

testing losser1
cuda time 60403.59765625
perf_counter time 60.351266200000005

testing losser3
cuda time 26321.48828125
perf_counter time 26.30165939999999

testing losser4
cuda time 24471.6875
perf_counter time 24.45189119999999

testing losser5
cuda time 23153.962890625
perf_counter time 23.135575399999993

pytorch 1.2

Performance Testing MS_SSIM

testing losser1
cuda time 134158.84375
perf_counter time 134.0433756

testing losser3
cuda time 62143.4140625
perf_counter time 62.103911400000015

testing losser4
cuda time 46854.25390625
perf_counter time 46.81785239999999

pytorch 1.1

Performance Testing MS_SSIM

testing losser1
cuda time 134115.96875
perf_counter time 134.0006031

testing losser3
cuda time 61760.56640625
perf_counter time 61.71994470000001

testing losser4
cuda time 52888.03125
perf_counter time 52.848280500000016

Test speed by yourself

cd ms_ssim_pytorch/_test_speed
python test_ssim_speed.py
or
python test_ms_ssim_speed.py

Other thing

Add parameter use_padding.
When set to True, the gaussian_filter behavior is the same as https://github.com/Po-Hsun-Su/pytorch-ssim.
This parameter is mainly used for MS-SSIM, because MS-SSIM needs to be downsampled.
When the input image is smaller than 176x176, this parameter needs to be set to True to ensure that MS-SSIM works normally. (when parameter weight and level are the default)

Require

Pytorch >= 1.1

if you want to test the code with animation. You also need to install some package.

pip install imageio imageio-ffmpeg opencv-python

Test code with animation

The test code is included in the ssim.py file, you can run the file directly to start the test.

git clone https://github.com/One-sixth/ms_ssim_pytorch
cd ms_ssim_pytorch
python ssim.py

Code Example.

import torch
import ssim


im = torch.randint(0, 255, (5, 3, 256, 256), dtype=torch.float, device='cuda')
img1 = im / 255
img2 = img1 * 0.5

losser = ssim.SSIM(data_range=1., channel=3).cuda()
loss = losser(img1, img2).mean()

losser2 = ssim.MS_SSIM(data_range=1., channel=3).cuda()
loss2 = losser2(img1, img2).mean()

print(loss.item())
print(loss2.item())