| 31 | |
| 32 | |
| 33 | def forward_flex(self, x): |
| 34 | b, c, h, w = x.shape |
| 35 | |
| 36 | pos_embed = self._resize_pos_embed( |
| 37 | self.pos_embed, h // self.patch_size[1], w // self.patch_size[0] |
| 38 | ) |
| 39 | |
| 40 | B = x.shape[0] |
| 41 | |
| 42 | if hasattr(self.patch_embed, "backbone"): |
| 43 | x = self.patch_embed.backbone(x) |
| 44 | if isinstance(x, (list, tuple)): |
| 45 | x = x[-1] # last feature if backbone outputs list/tuple of features |
| 46 | |
| 47 | x = self.patch_embed.proj(x).flatten(2).transpose(1, 2) |
| 48 | |
| 49 | if getattr(self, "dist_token", None) is not None: |
| 50 | cls_tokens = self.cls_token.expand( |
| 51 | B, -1, -1 |
| 52 | ) # stole cls_tokens impl from Phil Wang, thanks |
| 53 | dist_token = self.dist_token.expand(B, -1, -1) |
| 54 | x = torch.cat((cls_tokens, dist_token, x), dim=1) |
| 55 | else: |
| 56 | if self.no_embed_class: |
| 57 | x = x + pos_embed |
| 58 | cls_tokens = self.cls_token.expand( |
| 59 | B, -1, -1 |
| 60 | ) # stole cls_tokens impl from Phil Wang, thanks |
| 61 | x = torch.cat((cls_tokens, x), dim=1) |
| 62 | |
| 63 | if not self.no_embed_class: |
| 64 | x = x + pos_embed |
| 65 | x = self.pos_drop(x) |
| 66 | |
| 67 | for blk in self.blocks: |
| 68 | x = blk(x) |
| 69 | |
| 70 | x = self.norm(x) |
| 71 | |
| 72 | return x |
| 73 | |
| 74 | |
| 75 | def _make_vit_b16_backbone( |