1

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.

Gil Pinsky
  • 2,138
  • 1
  • 10
  • 14
Hao Zhou
  • 23
  • 4

1 Answers1

2

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.

jperezmartin
  • 414
  • 4
  • 19
Shai
  • 102,241
  • 35
  • 217
  • 344
  • Appreciate for your reply. This did solve my question. What if C subject to C[:, :, i, j] = torch.mean(torch.pow(A[:, :, i:i+80, j:j+80] - B, 2))? – Hao Zhou Oct 12 '20 at 09:20
  • @HaoZhou Ha! This is a nice exercise! (I'll probably give it to my students...). Here's a hint: `(a-b)^2 = a^2 - 2ab + b^2`. You already know how to handle the `a^2` and `b^2` parts. Think carefully what (very popular) operation can help you efficiently compute `ab` – Shai Oct 12 '20 at 09:57