MSSIM和L1loss的混合损失函数用于图像恢复

最新推荐文章于 2024-06-29 23:07:07 发布

群星闪耀

最新推荐文章于 2024-06-29 23:07:07 发布

阅读量1w

点赞数 8

分类专栏：深度学习学习记录 pytorch 文章标签：深度学习 python

本文链接：https://blog.csdn.net/weixin_40244676/article/details/119329980

版权

本文介绍了在图像恢复的深度学习任务中，如何结合MSSIM（多尺度结构相似性）和L1损失函数来提升恢复效果。文章详细阐述了MS_SSIM的计算过程，并提供了实现该混合损失函数的代码示例，强调了参数设置和输入图像的要求。

摘要由CSDN通过智能技术生成

背景

在图像恢复的深度神经网络当中。大部分的损失函数一般是L1 norm(MAE),L2 norm(MSE)。或者是SSIM,PSNR。
其中SSIM有了评价函数的变体MS_SSIM,具体的实现先参考文献[2], 代码部分再去参考[1]。
英伟达和MIT联合发表了一篇Loss Functions for Image Restoration with Neural Networks，这种损失函数将MS_SIM和L1损失函数结合在一起。
下面写一些具体如何实现

实现过程

MS_SSIM的实现过程已经在[1]中体现了，但是我推荐去原文[2]中把公式看一遍，对这个MS_SSIM的理解会深刻一些。
特别要注意的是，MS_SSIM中：
\begin{equation} \text{MS-SSIM}§ = l^\alpha _M§ \cdot \prod _{j = 1}^M cs^{\beta _j}_j§ \end{equation}

中计算亮度对比：\begin{equation} \text{} l^\alpha _M§ \end{equation}

只有最后M的，不像计算对比度因子的时候，要有对不同j={1，2，···，M}, 对比度因子C(x,y)求和。

实现过程如下，只要保存为文件名字,比如MSSIM_L1_loss.py。通过import导入class MS_SSIM_L1_LOSS，使用即可。

需要注意的是，看清楚输入图像的类型是cuda，tensor。还有图像的范围，我设置的是255，默认对灰度图像二维进行计算。如果是对彩色图像，有三通道，则需要把channel修改为3。

还有文献[3]里面给的alpha=0.16（论文里面的alpha是0.84，我这里是相反的），我设置为0.025，因为测试过了0.16效果不好

代码

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/8/2 17:09
# @Author  : Qilin Zhejiang University
import torch
import torch.nn as nn
import torch.nn.functional as F


class MS_SSIM_L1_LOSS(nn.Module):
    """
    Have to use cuda, otherwise the speed is too slow.
    Both the group and shape of input image should be attention on.
    I set 255 and 1 for gray image as default.
    """

    def __init__(self, gaussian_sigmas=[0.5, 1.0, 2.0, 4.0, 8.0],
                 data_range=255.0,
                 K=(0.01, 0.03),  # c1,c2
                 alpha=0.025,  # weight of ssim and l1 loss
                 compensation=200.0,  # final factor for total loss
                 cuda_dev=0,  # cuda device choice
                 channel=1):  # RGB image should set to 3 and Gray image should be set to 1
        super(MS_SSIM_L1_LOSS, self).__init__()
        self.channel = channel
        self.DR = data_range
        self.C1 = (K[0] * data_range) ** 2
        self.