Find elements surrounding an element in 2D tensor

vision
1

Hello

Suppose I have the following 2D tensor

 [[20, 29,  9,  7, 31, 35]],

        [[ 9, 23, 30,  5,  2, 31]],

        [[ 6, 27, 19,  5, 30,  7]],

        [[29,  4, 36,  2, 20, 27]],

        [[32,  2, 38, 10, 35,  1]],

        [[27, 40, 13, 10, 20,  3]],

        [[ 6, 20, 30,  7, 40, 39]]

I would like to create a new tensor, which will have the shape (…original_shape…,4)
And will contain the 4 surrounding neighbors of each element, with a special index to the border.
The directions will be “hard coded” (top,right,down,left)
For example,

new_tensor[0,0,:] = [[[
BORDER,29,BORDER,BORDER
]]]

I thought on 2 options:

  • Create surround function that will return the 4 surrounding elements (including BORDER in the right places) and run it on all tensor elements, and then build a new tensor from the results. But, as you may think, this is not the most efficient thing to do…

  • Use other tools, scatter / gather / others

Do you have any recommendation?

Thanks
Tankwell

2

Hi Tankwell!

Pad your input tensor with your special “border value” and then run conv2d()
with the appropriately constructed kernel:

>>> import torch
>>> print (torch.__version__)
1.12.0
>>>
>>> t = torch.tensor ([[20, 29,  9,  7, 31, 35],
...                    [ 9, 23, 30,  5,  2, 31],
...                    [ 6, 27, 19,  5, 30,  7],
...                    [29,  4, 36,  2, 20, 27],
...                    [32,  2, 38, 10, 35,  1],
...                    [27, 40, 13, 10, 20,  3],
...                    [ 6, 20, 30,  7, 40, 39]],
...                   dtype = torch.float)
>>>
>>> border_value = -99.0
>>>
>>> # build convolution kernel
>>> kernel = torch.zeros (4, 1, 3, 3)
>>> kernel[0, 0, 0, 1] = 1.0   # top
>>> kernel[1, 0, 1, 2] = 1.0   # right
>>> kernel[2, 0, 2, 1] = 1.0   # down
>>> kernel[3, 0, 1, 0] = 1.0   # left
>>>
>>> tpad = torch.nn.functional.pad (t, (1, 1, 1, 1), mode = 'constant', value = border_value)
>>> tsur = torch.nn.functional.conv2d (tpad.unsqueeze (0), kernel).permute (1, 2, 0)
>>>
>>> tsur.shape
torch.Size([7, 6, 4])
>>> tsur[0, 0, :]
tensor([-99.,  29.,   9., -99.])
>>> tsur
tensor([[[-99.,  29.,   9., -99.],
         [-99.,   9.,  23.,  20.],
         [-99.,   7.,  30.,  29.],
         [-99.,  31.,   5.,   9.],
         [-99.,  35.,   2.,   7.],
         [-99., -99.,  31.,  31.]],

        [[ 20.,  23.,   6., -99.],
         [ 29.,  30.,  27.,   9.],
         [  9.,   5.,  19.,  23.],
         [  7.,   2.,   5.,  30.],
         [ 31.,  31.,  30.,   5.],
         [ 35., -99.,   7.,   2.]],

        [[  9.,  27.,  29., -99.],
         [ 23.,  19.,   4.,   6.],
         [ 30.,   5.,  36.,  27.],
         [  5.,  30.,   2.,  19.],
         [  2.,   7.,  20.,   5.],
         [ 31., -99.,  27.,  30.]],

        [[  6.,   4.,  32., -99.],
         [ 27.,  36.,   2.,  29.],
         [ 19.,   2.,  38.,   4.],
         [  5.,  20.,  10.,  36.],
         [ 30.,  27.,  35.,   2.],
         [  7., -99.,   1.,  20.]],

        [[ 29.,   2.,  27., -99.],
         [  4.,  38.,  40.,  32.],
         [ 36.,  10.,  13.,   2.],
         [  2.,  35.,  10.,  38.],
         [ 20.,   1.,  20.,  10.],
         [ 27., -99.,   3.,  35.]],

        [[ 32.,  40.,   6., -99.],
         [  2.,  13.,  20.,  27.],
         [ 38.,  10.,  30.,  40.],
         [ 10.,  20.,   7.,  13.],
         [ 35.,   3.,  40.,  10.],
         [  1., -99.,  39.,  20.]],

        [[ 27.,  20., -99., -99.],
         [ 40.,  30., -99.,   6.],
         [ 13.,   7., -99.,  20.],
         [ 10.,  40., -99.,  30.],
         [ 20.,  39., -99.,   7.],
         [  3., -99., -99.,  40.]]])

(If you can live with 0.0 as your special border value, you can avoid
the cost in time and memory of constructing tpad by using Conv2d with
padding_mode = 'zeros'.)

Best.

K. Frank

3

Hey,
Tha’ts relly helped!

Actually I have to select some of the centers and not all of them.
I have 2 ways to do that:

  1. Add them to the surround tensor above by 5 nonzero elements [add (1,1)] in the kernel, and looking for that value later
  2. Use the original 2D tensor to find [for each sample] the i,j location of the argument passed (the value we are looking for). After, searching in that location on surr

The thing is that I have to locate a value in a 2D tensor on each sample in the batch…seems to be not so efficient