[pytorch] 计算图像的一阶导 / 梯度 / gradient
在图像转换任务中常见的total variation loss(tvloss,总变分,一般作为平滑的规则化项)需要对图像的梯度求平方和。
style-transfer系的github项目,tvloss求法如下:
class TVLoss(torch.nn.Module):
def __init__(self):
super(TVLoss,self).__init__()
def forward(self,x):
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