How can I replicate this tensor.unfold operation with F.unfold?

I need the dilation parameter, and tensor.unfold lacks it.

import torch


x = torch.randn(1,32,16,16)

kernel_size=(3,3)
dilation = 1
stride=(1,1)

x_out = x.unfold(2, kernel_size[0], stride[0]).unfold(3, kernel_size[1], stride[1])
print(x_out.shape)


x_out = torch.nn.functional.unfold(x, kernel_size=kernel_size, dilation=dilation, padding=0, stride=stride)
print(x_out.shape)

nn.MaxPool2d has a dilation value of 1, and I’m using the above to make AvgPool2d match MaxPool2d’s parameters.

Edit: Looks like I don’t need it after all.

1 Like