[Refactor] Refactor _prepare_pos_embed in ViT (open-mmlab#1656)

* deal with cls_token * Update implement --------- Co-authored-by: mzr1996 <[email protected]>
ASHORE1225 · Jun 20, 2023 · 70ff2ab · 70ff2ab
1 parent d4a6dfa
commit 70ff2ab
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/mmpretrain/models/backbones/vision_transformer.py b/mmpretrain/models/backbones/vision_transformer.py
@@ -305,6 +305,7 @@ def __init__(self,
         self.out_type = out_type
 
         # Set cls token
+        self.with_cls_token = with_cls_token
         if with_cls_token:
             self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims))
         elif out_type != 'cls_token':
@@ -404,6 +405,11 @@ def _prepare_pos_embed(self, state_dict, prefix, *args, **kwargs):
                 int(np.sqrt(ckpt_pos_embed_shape[1] - self.num_extra_tokens)))
             pos_embed_shape = self.patch_embed.init_out_size
 
+            if (not self.with_cls_token and ckpt_pos_embed_shape[1]
+                    == self.pos_embed.shape[1] + 1):
+                # Remove cls token from state dict if it's not used.
+                state_dict[name] = state_dict[name][:, 1:]
+
             state_dict[name] = resize_pos_embed(state_dict[name],
                                                 ckpt_pos_embed_shape,
                                                 pos_embed_shape,