Pytorch: Need help implementing a custom sliding window

Viewed 224

Lets say I have two matrices A and B, both are 10x10. Now I want a sliding window, 3x3 to go through A and B and:

  1. calculate the following result:

enter image description here

  1. then take this result matrix, sum it up, and save it at location e inside of a new matrix.
1 Answers

Given you are applying a weighted sum (i.e. a convolution), your kernel is equivalent to the following:

enter image description here

Which can be straightforwardly implemented with torch.nn.functional.conv2d:

import torch.nn.functional as F

kernel = torch.Tensor([[-1,-2,-3],
                       [-4,sum([1,2,3,4,6,7,8,9]), -6],
                       [-7,-8,-9]])
kernel = kernel.reshape(1,1,3,3)

A = torch.randn(1,5,5)
B = torch.randn(1,5,5)
inputs = torch.stack([A,B])
F.conv2d(inputs, kernel, padding=1)
Related