Down-resing a 3D tensor

1

I have a tensor of shape [5sx,5sy,5*sz] and I’d like to create a tensor with shape [sx,sy,sz] where the value of each element:

new_tensor[x,y,z] = mean(old_tensor[5x:5x+5,5y:5y+5,5z:5z+5])

Preferably in one instruction that’s fast (and followed by autograd)

Any ideas?
Thanks,
Ian

2

kernel_size=[5,5,5]
torch.nn.Conv3d(in_channels=1,out_channels=1,kernel_size=kernel_size, stride=kernel_size, padding=0)

3

Hi Ian!

You might prefer to use AvgPool3d. You can view it as a special
case of Conv3d:

>>> import torch
>>> torch.__version__
'1.10.2'
>>> _ = torch.manual_seed (2022)
>>> t = torch.randn (5, 10, 15, requires_grad = True)
>>> torch.nn.AvgPool3d (5) (t.unsqueeze (0).unsqueeze (0)).squeeze (0).squeeze (0)
tensor([[[-0.2630, -0.0570,  0.1348],
         [ 0.2163,  0.0359,  0.0404]]], grad_fn=<SqueezeBackward1>)
>>> conv = torch.nn.Conv3d (1, 1, 5, stride = 5, bias = False)
>>> conv.weight.requires_grad = False
>>> _ = conv.weight.copy_ (torch.ones (1, 1, 5, 5, 5) / 125)
>>> conv (t.unsqueeze (0).unsqueeze (0)).squeeze (0).squeeze (0)
tensor([[[-0.2630, -0.0570,  0.1348],
         [ 0.2163,  0.0359,  0.0404]]], grad_fn=<SqueezeBackward1>)

I haven’t done any timings, but I would expect AvgPool3d to be
modestly faster.

And stylistically, AvgPool3d – at least for me – more clearly “says
what you mean.”

Best.

K. Frank

1 Like