pytorch | 计算图像的一阶导/梯度/gradient

[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]),2).sum()
        return h_tv/count_h + w_tv/count_w

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

根据上面可以看出，这个求法避免了 （x[:,:,1:,:]-x[:,:,:h_x-1,:]）与（x[:,:,:,1:]-x[:,:,:,:w_x-1]）维度不一样的问题。
那么如果想要求得梯度作为变量怎么办咧？
我参考了numpy源码的写法，尝试利用torch.nn.functional.pad()这个函数实现了一个一阶导的版本，其中dx取值为2个像素：

import torch
import torch.nn.functional as F
def gradient_1order(x,h_x=None,w_x=None):
    if h_x is None and w_x is None:
        h_x = x.size()[2]
        w_x = x.size()[3]
    r = F.pad(x, (0, 1, 0, 0))[:, :, :, 1:]
    l = F.pad(x, (1, 0, 0, 0))[:, :, :, :w_x]
    t = F.pad(x, (0, 0, 1, 0))[:, :, :h_x, :]
    b = F.pad(x, (0, 0, 0, 1))[:, :, 1:, :]
    xgrad = torch.pow(torch.pow((r - l) * 0.5, 2) + torch.pow((t - b) * 0.5, 2), 0.5)
    return xgrad

测试

import torch
npim = np.array(im,dtype=np.float32)
import numpy as np
import PIL.Image as Image

def gradient_numpy(img):
    # input : Image/nparray 0~255
    changeflag = isinstance(img, Image.Image)
    if changeflag:
        img = np.array(img)
    dx,dy = np.gradient(img/255,edge_order=1)
    img = (np.sqrt(dx**2 + dy**2)*255).astype(np.uint8)
    if changeflag:
        img = Image.fromarray(img)
    return img

im = Image.open('text.jpg').convert('L').crop([1,1,200,200])
# fig 1
im.show()
# fig 2 
gradient_numpy(im).show()

npim = np.array(im,dtype=np.float32)
tim = torch.from_numpy(npim).unsqueeze_(0).unsqueeze_(0)
gradient_1order(tim)
npgrad = t.squeeze(0).squeeze(0).data.clamp(0,255).numpy()
# fig 3
Image.fromarray(npgrad.astype('uint8')).show()

---

从效果上看是没问题的，但是遇到个我不解的情况：图像pair之间如果用l2loss或l1loss都可以用作生成网络的loss，但为什么我用图像的梯度pair做l2loss或l1loss就不可以呢？反向传播的时候生成网络直接梯度爆炸了。。。
并不是只用了这个梯度l2loss，而是不管这个loss权重多小，学习率多低，都会爆炸，不是很明白。。。。