pytorch 计算图像 x,y 梯度，模仿 tf.image.

tf.image.image_gradients：https://www.tensorflow.org/api_docs/python/tf/image/image_gradients

pytroch use pad to implement

import torch.nn.functional as F
import torch
import numpy as np


def gradient(x):
    # tf.image.image_gradients(image)
    h_x = x.size()[-2]
    w_x = x.size()[-1]
    # gradient step=1
    l = x
    r = F.pad(x, [0, 1, 0, 0])[:, :, :, 1:]
    t = x
    b = F.pad(x, [0, 0, 0, 1])[:, :, 1:, :]

    dx, dy = torch.abs(r - l), torch.abs(b - t)
    # dx will always have zeros in the last column, r-l
    # dy will always have zeros in the last row,    b-t
    dx[:, :, :, -1] = 0
    dy[:, :, -1, :] = 0

    return dx, dy


# fake depth image
a = np.arange(1, 26).reshape((1, 1, 5, 5))
a = torch.from_numpy(a)
print('a')
print(a)
# print('pad left')
# b = F.pad(a, [1, 0, 0, 0])  # use 0 pad
# print(b)

dx, dy = gradient(a)

print('dx')
print(dx)
print('dy')
print(dy)

a
tensor([[[[ 1,  2,  3,  4,  5],
          [ 6,  7,  8,  9, 10],
          [11, 12, 13, 14, 15],
          [16, 17, 18, 19, 20],
          [21, 22, 23, 24, 25]]]])
dx
tensor([[[[1, 1, 1, 1, 0],
          [1, 1, 1, 1, 0],
          [1, 1, 1, 1, 0],
          [1, 1, 1, 1, 0],
          [1, 1, 1, 1, 0]]]])
dy
tensor([[[[5, 5, 5, 5, 5],
          [5, 5, 5, 5, 5],
          [5, 5, 5, 5, 5],
          [5, 5, 5, 5, 5],
          [0, 0, 0, 0, 0]]]])