fastest way to calculate edges (derivatives) of a big torch tensor
Given a tensor with shape (b,c,h,w), I want to extract edges of the spatial data, that is, calculate x, y direction derivatives of the (h,w) and calculate the magnitude I=sqrt(|x_amplitude|^2+|y_amplitude|^2)

My current implementation is as followed
row_mat = np.asarray([[0, 0, 0], [1, 0, -1], [0, 0, 0]])
col_mat = row_mat.T
row_mat = row_mat[None, None, :, :]  # expand dim to convolve with tensor (batch,channel,width,height)
col_mat = col_mat[None, None, :, :]  # expand dim to convolve with tensor (batch,channel,width,height)

def derivative(batch: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;
    uses convolution to perform x and y derivatives
    :param batch: input tensor batch
    :return: image derivative magnitudes
    &quot;&quot;&quot;
    x_amplitude = ndimage.convolve(batch, row_mat)
    y_amplitude = ndimage.convolve(batch, col_mat)
    magnitude = np.sqrt(np.abs(x_amplitude) ** 2 + np.abs(y_amplitude) ** 2)
    return torch.tensor(magnitude)

I was wondering if there's a faster way, as this approach actually convolves using the definition of a derivative, so there might be downsides to that.

PS. to test this you can use the tensor torch.randn(1000,128,28,28), as these are the dimension I'm dealing with

For this specific operation you might be able to speed things up a bit by doing it &quot;manually&quot;:
import torch.nn.functional as nnf

def derivative(batch: torch.Tensor) -&gt; torch.Tensor:
  # pad batch
  x = nnf.pad(batch, (1, 1, 1, 1), mode='reflect')
  dx2 = (x[..., 1:-2, :-2] - x[..., 1:-2, 2:])**2
  dy2 = (x[..., :-2, 1:-2] - x[..., 2:, 1:-2])**2
  mag = torch.sqrt(dx2 + dy2)
  return mag

