MS-SSIM-PyTorch是一个利用PyTorch JIT技术优化过的高效计算结构相似性指标(SSIM)及其多尺度版本(MS-SSIM)的代码库,适用于图像和视频的质量评估。
这段代码是从已有的代码基础上修改而来的,并且部分代码已经经过优化以提高速度、减少VRAM占用并兼容pytorch jit。
还有一个动态频道版本可供选择,这个版本使用起来更加方便,但性能损失很小。
感谢贡献者找到并修复了ms_ssim在反向传播时会导致梯度nan的问题。
如果您正在使用的是pytorch 1.2,请注意不要在训练循环中创建和销毁该jit模块(其他jit模块也可能出现这种情况),这可能会导致内存泄漏。我已经测试过,在pytorch 1.6版本上没有这个问题。
我还研究了piqa库,这使得我实现的ssim和ms-ssim的速度比以前更快了一些。
加速效果仅在GPU上进行了验证。
losser1是268fc76, losser2是881d210, losser3是5caf547, losser4是1c2f14a, loser5是abaf398,其中最新的版本为abaf398。
在pytorch 1.7中进行了相应的优化。