| 173 | return x_cls, attn |
| 174 | |
| 175 | class Patch_Block(nn.Module): |
| 176 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., |
| 177 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): |
| 178 | super().__init__() |
| 179 | self.norm1 = norm_layer(dim) |
| 180 | self.attn = Patch_Attention( |
| 181 | dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) |
| 182 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 183 | self.norm2 = norm_layer(dim) |
| 184 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 185 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 186 | |
| 187 | def forward(self, x_cls, x, return_attention=False): |
| 188 | u = torch.cat((x_cls,x),dim=1) |
| 189 | y, attn = self.attn(self.norm1(u)) |
| 190 | if return_attention: |
| 191 | return attn |
| 192 | x = x + self.drop_path(y) |
| 193 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 194 | return x |
| 195 | |
| 196 | class Attention(nn.Module): |
| 197 | def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): |
nothing calls this directly
no outgoing calls
no test coverage detected