pytorch:自定义损失函数total variation loss

最新推荐文章于 2024-08-12 18:16:40 发布

lloyd_ljj

最新推荐文章于 2024-08-12 18:16:40 发布

阅读量1w

点赞数 6

文章标签：深度学习 pytorch loss 损失函数 total variation

本文链接：https://blog.csdn.net/qq_34622844/article/details/88846411

版权

DeepLearningAndPytorch 专栏收录该内容

3 篇文章 0 订阅

订阅专栏

Total variation loss常被用在损失函数里的正则项，可以起到平滑图像，去除鬼影，消除噪声的作用
TVloss的表达式如下所示：
total variation loss 表达式
下面是在pytorch中的代码：

import torch
import torch.nn as nn
from torch.autograd import Variable

class TVLoss(nn.Module):
    def __init__(self,TVLoss_weight=1):
        super(TVLoss,self).__init__()
        self.TVLoss_weight = TVLoss_weight

    def forward(self,x):
        batch_size = x.size()[0]
        h_x = x.size()[2]
        w_x = x.size()[3]
        count_h =  (x.size()[2]-1) * x.size()[3]
        count_w = x.size()[2] * (x.size()[3] - 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.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size

def 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)

if __name__ == '__main__':
    main()