| 360 | return model |
| 361 | |
| 362 | class SelfPatchHead(nn.Module): |
| 363 | def __init__(self, in_dim, num_heads, k_num): |
| 364 | super().__init__() |
| 365 | self.cls_token = nn.Parameter(torch.zeros(1, 1, in_dim)) |
| 366 | self.cls_blocks = nn.ModuleList([ |
| 367 | LayerScale_Block_CA( |
| 368 | dim=in_dim, num_heads=num_heads, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, |
| 369 | drop=0.0, attn_drop=0.0, drop_path=0.0, norm_layer=partial(nn.LayerNorm, eps=1e-6), |
| 370 | act_layer=nn.GELU, Attention_block=Class_Attention, |
| 371 | Mlp_block=Mlp) |
| 372 | for i in range(2)]) |
| 373 | trunc_normal_(self.cls_token, std=.02) |
| 374 | self.norm = partial(nn.LayerNorm, eps=1e-6)(in_dim) |
| 375 | |
| 376 | self.apply(self._init_weights) |
| 377 | self.k_num = k_num |
| 378 | self.k_size = 3 |
| 379 | self.loc224 = self.get_local_index(196, self.k_size) |
| 380 | self.loc96 = self.get_local_index(36, self.k_size) |
| 381 | self.embed_dim = in_dim |
| 382 | |
| 383 | def _init_weights(self, m): |
| 384 | if isinstance(m, nn.Linear): |
| 385 | trunc_normal_(m.weight, std=.02) |
| 386 | if isinstance(m, nn.Linear) and m.bias is not None: |
| 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): |