(self, x_cls, x, return_attention=False)
| 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.): |