| 387 | nn.init.constant_(m.bias, 0) |
| 388 | |
| 389 | def forward(self, x, loc=False): |
| 390 | cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) |
| 391 | if loc: |
| 392 | k_size = self.k_size |
| 393 | if x.shape[1] == 196: |
| 394 | local_idx = self.loc224 |
| 395 | elif x.shape[1] == 36: |
| 396 | if self.k_size == 14: |
| 397 | k_size = 6 |
| 398 | local_idx = self.loc96 |
| 399 | else: |
| 400 | print(x.shape) |
| 401 | assert(False) |
| 402 | |
| 403 | x_norm = nn.functional.normalize(x, dim=-1) |
| 404 | sim_matrix = x_norm[:,local_idx] @ x_norm.unsqueeze(2).transpose(-2,-1) |
| 405 | top_idx = sim_matrix.squeeze().topk(k=self.k_num,dim=-1)[1].view(-1,self.k_num,1) |
| 406 | |
| 407 | x_loc = x[:,local_idx].view(-1,k_size**2-1,self.embed_dim) |
| 408 | x_loc = torch.gather(x_loc, 1, top_idx.expand(-1, -1, self.embed_dim)) |
| 409 | for i, blk in enumerate(self.cls_blocks): |
| 410 | if i == 0: |
| 411 | glo_tokens = blk(x, cls_tokens) |
| 412 | loc_tokens = blk(x_loc, cls_tokens.repeat(x.shape[1],1,1)) |
| 413 | else: |
| 414 | glo_tokens = blk(x, glo_tokens) |
| 415 | loc_tokens = blk(x_loc, loc_tokens) |
| 416 | loc_tokens = loc_tokens.view(x.shape) |
| 417 | x = self.norm(torch.cat([glo_tokens, loc_tokens], dim=1)) |
| 418 | else: |
| 419 | for i, blk in enumerate(self.cls_blocks): |
| 420 | cls_tokens = blk(x, cls_tokens) |
| 421 | x = self.norm(torch.cat([cls_tokens, x], dim=1)) |
| 422 | |
| 423 | return x |
| 424 | |
| 425 | @staticmethod |
| 426 | def get_local_index(N_patches, k_size): |