(self, in_dim, num_heads, k_num)
| 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): |
nothing calls this directly
no test coverage detected