How to handle variant image sizes? Thanks

[junyongyou]

I have a question about variant image sizes.

If we have images with different sizes (actually happens often, if no resizing is used). Let's say imge1 has 256 patches, and image2 has 512 patches. For this question, I would guess self.pos_embedding is defined as a sufficient big, e.g.,
self.pos_embedding = nn.Parameter(torch.randn(1, 10000, dim)), and then when using it, we may use
num_patches = x.shape[1] x += self.pos_embedding(:, num_patches + 1, :).
But I am not quite sure if this approach works. Could you please advise?

[lucidrains]

@junyongyou Hi! I just released a new version of the library that will allow you to do this, provided when passing image_size in the constructor, you pass the maximum image size

import torch
from vit_pytorch import ViT

model = ViT(
	image_size = 512,
	patch_size = 32,
	num_classes = 1000,
	dim = 1024,
	depth = 6,
	heads = 8,
	mlp_dim = 2048,
	dropout = 0.1,
	emb_dropout = 0.1
)

x = torch.randn(1, 3, 256, 256)
y = torch.randn(1, 3, 512, 512)
model(x), model(y)

the other way to do this, is to embed your smaller image into the largest image size, then pass the appropriate mask to block out attention to the non-image regions

ex.

mask = torch.tensor([
  [1, 1, 0, 0]
  [1, 1, 0, 0]
  [0, 0, 0, 0]
  [0, 0, 0, 0]
]).bool()

x = F.pad(x, (0, 256, 0, 256), 0)
model(x, mask = mask)

[lucidrains]

Just throwing some caution though; you will need a sufficient number of patches for attention to work well, as attention relies on the comparison between the representative tokens of each patch

[lucidrains]

In other words, if you pass in a 32x32 image (1 token), nothing will get learned in a transformer