Spaces:
Runtime error
Runtime error
File size: 4,682 Bytes
3424266 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
# Copyright 2024 EPFL and Apple Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# --------------------------------------------------------
# Based on UniLM / BEiT code base
# https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------
import re
import torch
def interpolate_pos_embed_vit(model, checkpoint_model):
if 'pos_embed' in checkpoint_model:
pos_embed_checkpoint = checkpoint_model['pos_embed']
embedding_size = pos_embed_checkpoint.shape[-1]
num_patches = model.patch_embed.num_patches
num_extra_tokens = model.pos_embed.shape[-2] - num_patches
# height (== width) for the checkpoint position embedding
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
# height (== width) for the new position embedding
new_size = int(num_patches ** 0.5)
# class_token and dist_token are kept unchanged
if orig_size != new_size:
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
# only the position tokens are interpolated
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
pos_tokens = torch.nn.functional.interpolate(
pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
checkpoint_model['pos_embed'] = new_pos_embed
def interpolate_pos_embed_multimae(model, checkpoint_model):
pattern = "input_adapters\.(.*)\.pos_emb"
matched_keys = [k for k in checkpoint_model if bool(re.match(pattern, k))]
for key in matched_keys:
domain = re.match(pattern, key).group(1) # group(0) is entire matched regex
if getattr(model.input_adapters, domain, None) is not None:
pos_embed_checkpoint = checkpoint_model[key]
_, _, orig_H, orig_W = pos_embed_checkpoint.shape
_, _, new_H, new_W = getattr(model.input_adapters, domain).pos_emb.shape
if (orig_H != new_H) or (orig_W != new_W):
print(f"Key {key}: Position interpolate from {orig_H}x{orig_W} to {new_H}x{new_W}")
pos_embed_checkpoint = torch.nn.functional.interpolate(
pos_embed_checkpoint, size=(new_H, new_W), mode='bicubic', align_corners=False)
checkpoint_model[key] = pos_embed_checkpoint
def interpolate_rgb_pos_emb_fm(model, checkpoint_model):
if 'encoder_embeddings.rgb.pos_emb' in checkpoint_model:
pos_embed_checkpoint = checkpoint_model['encoder_embeddings.rgb.pos_emb']
embedding_size = pos_embed_checkpoint.shape[-1]
num_patches = model.encoder_embeddings.rgb.num_patches
num_extra_tokens = 0
# height (== width) for the checkpoint position embedding
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
# height (== width) for the new position embedding
new_size = int(num_patches ** 0.5)
# class_token and dist_token are kept unchanged
if orig_size != new_size:
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
# only the position tokens are interpolated
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
pos_tokens = torch.nn.functional.interpolate(
pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
checkpoint_model['encoder_embeddings.rgb.pos_emb'] = new_pos_embed
|