| 320 | """ |
| 321 | |
| 322 | def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768): |
| 323 | super().__init__() |
| 324 | assert isinstance(backbone, nn.Module) |
| 325 | img_size = to_2tuple(img_size) |
| 326 | self.img_size = img_size |
| 327 | self.backbone = backbone |
| 328 | if feature_size is None: |
| 329 | with torch.no_grad(): |
| 330 | # FIXME this is hacky, but most reliable way of determining the exact dim of the output feature |
| 331 | # map for all networks, the feature metadata has reliable channel and stride info, but using |
| 332 | # stride to calc feature dim requires info about padding of each stage that isn't captured. |
| 333 | training = backbone.training |
| 334 | if training: |
| 335 | backbone.eval() |
| 336 | o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1] |
| 337 | feature_size = o.shape[-2:] |
| 338 | feature_dim = o.shape[1] |
| 339 | backbone.train(training) |
| 340 | else: |
| 341 | feature_size = to_2tuple(feature_size) |
| 342 | feature_dim = self.backbone.feature_info.channels()[-1] |
| 343 | self.num_patches = feature_size[0] * feature_size[1] |
| 344 | self.proj = nn.Linear(feature_dim, embed_dim) |
| 345 | |
| 346 | def forward(self, x): |
| 347 | x = self.backbone(x)[-1] |