图像去燥——TV Loss

1. 背景介绍

        在图像复原过程中，图像上的一点点噪声可能就会对复原的结果产生非常大的影响，因为很多复原算法都会放大噪声。这时候我们就需要在最优化问题的模型中添加一些正则项来保持图像的光滑性，TV loss是常用的一种正则项（注意是正则项，配合其他loss一起使用，约束噪声）。图片中相邻像素值的差异可以通过降低TV loss来一定程度上解决。比如降噪，对抗checkerboard等等。

2. 公式表达 

      即：求每一个像素和横向下一个像素的差的平方，加上纵向下一个像素的差的平方。然后开β/2次根。一般β的默认值为2，所以可以不用考虑，β/2。

      TVLoss促进了生成的图像中的空间平滑性:当β < 1时，会出现下图左侧的小点点的artifact；当β > 1时，图像中小点点会被消除，但是代价就是图像的清晰度。效果图如下：

 

3. 代码实现 

import torch
from torch.autograd import Variable

class TVLoss(torch.nn.Module):
    """
    TV loss
    """

    def __init__(self, weight=1):
        super(TVLoss, self).__init__()
        self.weight = weight

    def forward(self, x):
        batch_size = x.size()[0]
        h_x = x.size()[2]
        w_x = x.size()[3]
        count_h = self._tensor_size(x[:,:,1:,:])
        count_w = self._tensor_size(x[:,:,:,1:])
        h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum()
        w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum()
        return self.weight*2*(h_tv/count_h+w_tv/count_w)/batch_size

    def _tensor_size(self, t):
        return t.size()[1]*t.size()[2]*t.size()[3]

if __name__ == "__main__":
    x = Variable(torch.FloatTensor([[[1, 2, 3], [2, 3, 4], [3, 4, 5]], [[1, 2, 3], [2, 3, 4], [3, 4, 5]]]).view(1, 2, 3, 3),requires_grad=True)
    addition = TVLoss()
    z = addition(x)
    print(z)
    z.backward()

  运行结果：tensor(4., grad_fn=<DivBackward0>)