CODEWITHSUNDEEP
pytorch
Hao Zhou
Hao Zhou

Reputation: 23

element-wise operation in pytorch

I have two Tensors A and B, A.shape is (b,c,100,100), B.shape is (b,c,80,80), how can I get tensor C with shape (b,c,21,21) subject to C[:, :, i, j] = torch.mean(A[:, :, i:i+80, j:j+80] - B)? I wonder whether there's an efficient way to solve this? Thanks very much.

Upvotes: 1

Views: 2126

Answers (1)

Shai
Shai

Reputation: 114786

You should use an average pool to compute the sliding window mean operation.
It is easy to see that:

mean(A[..., i:i+80, j:j+80] - B) = mean(A[..., i:i+80, j:j+80]) - mean(B)

Using avg_pool2d:

import torch.nn.functional as nnf

C = nnf.avg_pool2d(A, kernel_size=80, stride=1, padding=0) - torch.mean(B, dim=(2,3), keepdim=True)

If you are looking for a more general way of performing sliding window operations in PyTorch, you should look at fold and unfold.

Upvotes: 2

Related Questions