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

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

下面是在pytorch中的代码：

import torch
import torch.nn as nn
from torch.autograd import Variableclass TVLoss(nn.Module):def __init__(self,TVLoss_weight=1):super(TVLoss,self).__init__()self.TVLoss_weight = TVLoss_weightdef 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_sizedef 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()