Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

interpolate positional embeddings for pretrained image tower #41

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

interpolate positional embeddings for pretrained image tower #41

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions src/open_clip/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -930,19 +930,25 @@ def trace_model(model, batch_size=256, device=torch.device('cpu')):
return model


def resize_eva_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
def resize_eva_pos_embed(state_dict, model, interpolation: str = 'bicubic'):
# interpolate position embedding
prepending = 'module.' if next(iter(state_dict)).startswith('module') else ''

visual_tower = (
'visual.' if any('visual' in key for key in state_dict.keys()) else ''
)
# Check if the position embedding is in the state_dict
pos_embed_key = f'{prepending}visual.pos_embed'
pos_embed_key = f'{prepending}{visual_tower}pos_embed'
if pos_embed_key in state_dict:
pos_embed_checkpoint = state_dict[pos_embed_key]
embedding_size = pos_embed_checkpoint.shape[-1]
num_patches = (
model.module.visual.patch_embed.num_patches
if hasattr(model, 'module')
else model.visual.patch_embed.num_patches
else (
model.visual.patch_embed.num_patches
if hasattr(model, 'visual')
else model.patch_embed.num_patches
)
)
num_extra_tokens = model.state_dict()[pos_embed_key].shape[-2] - num_patches
# height (== width) for the checkpoint position embedding
Expand Down Expand Up @@ -972,12 +978,16 @@ def resize_eva_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_
state_dict[pos_embed_key] = new_pos_embed

# Resize the patch embedding projection
patch_embed_key = f'{prepending}visual.patch_embed.proj.weight'
patch_embed_key = f'{prepending}{visual_tower}patch_embed.proj.weight'
patch_embed_proj = state_dict[patch_embed_key]
patch_size = (
model.module.visual.patch_embed.patch_size
if hasattr(model, 'module')
else model.visual.patch_embed.patch_size
else (
model.visual.patch_embed.patch_size
if hasattr(model, 'visual')
else model.patch_embed.patch_size
)
)

state_dict[patch_embed_key] = torch.nn.functional.interpolate(
Expand Down