全变分TV损失计算

最新推荐文章于 2024-06-16 20:31:37 发布

天天都在摸鱼的乐乐

最新推荐文章于 2024-06-16 20:31:37 发布

阅读量501

点赞数

文章标签： python 深度学习 pytorch

本文链接：https://blog.csdn.net/qq_46401672/article/details/130410049

版权

全变分TV损失计算：
python实现：
在这里插入图片描述
其中，C、H、W分别表示图像的通道、高度和宽度。∇x和∇y分别为水平梯度运算和垂直梯度运算

class L_TV(nn.Module):
    def __init__(self,TVLoss_weight=1):
        super(L_TV,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