0

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) -> torch.Tensor:
    """
    uses convolution to perform x and y derivatives
    :param batch: input tensor batch
    :return: image derivative magnitudes
    """
    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

Shai
  • 111,146
  • 38
  • 238
  • 371
Hadar
  • 658
  • 4
  • 17

1 Answers1

1

For this specific operation you might be able to speed things up a bit by doing it "manually":

import torch.nn.functional as nnf

def derivative(batch: torch.Tensor) -> 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
Shai
  • 111,146
  • 38
  • 238
  • 371