ViT for resolution beyond 224x224 support

When the resolution changes, the size of position embedding of ViTModel also changes, which makes from_pretrained method not working.

So, how can I use ViT with a different resolution like 64x64?

[NielsRogge]

I think you first need to load the state_dict of the original model, like so:

from transformers import ViTModel

model = ViTModel.from_pretrained('google/vit-base-patch16-224') # load pretrained model
state_dict = model.state_dict()

Then, initialize a new ViTModel with custom image_size, update the position embeddings of the state_dict and load the new model with that state_dict:

from transformers import ViTConfig

config = ViTConfig.from_pretrained('google/vit-base-patch16-224', image_size=64)
# new model with custom image_size
model = ViTModel(config=config)

# update state_dict
new_state_dict = state_dict.copy()
old_posemb = new_state_dict['embeddings.position_embeddings']
if model.embeddings.position_embeddings.shape != old_posemb.shape: # need to resize the position embedding by interpolate
    new_posemb = resize_pos_embed(old_posemb, model.embeddings.position_embeddings) # use PyTorch function linked above
    new_state_dict['embeddings.position_embeddings'] = new_posemb

# equip new model with state_dict
model.load_state_dict(new_state_dict)

[NielsRogge]

One would need to interpolate the pre-trained position embeddings. You can see how this is done in the original implementation here.

You can find a PyTorch implementation of that here (taken from the T2T-ViT implementation), where they show how you can go from 224 to 384. The pre-trained position embeddings are of shape (1, 197, 768) - there are 196 “positions” in an image of 224x224 with a patch size of 16x16 as (224/16)^2 = 196 and we add 1 for the [CLS] token - and suppose you want to fine-tune at resolution of 64x64 with a patch size of 8, then the number of position embeddings is (64/8)^2 + 1 = 65. In that case, the position embeddings during fine-tuning are of shape (1, 65, 768), and you can use that function to map the pre-trained position embeddings from shape (1, 197, 768) to (1, 65, 768).