| 73 | # DeiT: https://github.com/facebookresearch/deit |
| 74 | # -------------------------------------------------------- |
| 75 | def interpolate_pos_embed(model, checkpoint_model): |
| 76 | if 'pos_embed' in checkpoint_model: |
| 77 | pos_embed_checkpoint = checkpoint_model['pos_embed'] |
| 78 | embedding_size = pos_embed_checkpoint.shape[-1] |
| 79 | num_patches = model.patch_embed.num_patches |
| 80 | num_extra_tokens = model.pos_embed.shape[-2] - num_patches |
| 81 | # height (== width) for the checkpoint position embedding |
| 82 | orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5) |
| 83 | # height (== width) for the new position embedding |
| 84 | new_size = int(num_patches ** 0.5) |
| 85 | # class_token and dist_token are kept unchanged |
| 86 | if orig_size != new_size: |
| 87 | print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size)) |
| 88 | extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens] |
| 89 | # only the position tokens are interpolated |
| 90 | pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:] |
| 91 | pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2) |
| 92 | pos_tokens = torch.nn.functional.interpolate( |
| 93 | pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False) |
| 94 | pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2) |
| 95 | new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1) |
| 96 | checkpoint_model['pos_embed'] = new_pos_embed |