Can we divide up input images into patches?

4

This code should work:

import torch

x = torch.arange(103*152*128).view(103, 152, 128)
print(x)

out = x.unfold(0, 16, 16).unfold(1, 16, 16)
print(out.shape)
# torch.Size([6, 9, 128, 16, 16])
out = out.permute(0, 1, 3, 4, 2).contiguous().view(-1, 16, 16, 128)
print(out.shape)
# torch.Size([54, 16, 16, 128])
print(out)

and also this post might be helpful, too.