| 189 | |
| 190 | |
| 191 | class Block(nn.Module): |
| 192 | |
| 193 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., |
| 194 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_gpsa=True, **kwargs): |
| 195 | super().__init__() |
| 196 | self.norm1 = norm_layer(dim) |
| 197 | self.use_gpsa = use_gpsa |
| 198 | if self.use_gpsa: |
| 199 | self.attn = GPSA( |
| 200 | dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, **kwargs) |
| 201 | else: |
| 202 | self.attn = MHSA(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) |
| 203 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 204 | self.norm2 = norm_layer(dim) |
| 205 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 206 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 207 | |
| 208 | def forward(self, x): |
| 209 | x = x + self.drop_path(self.attn(self.norm1(x))) |
| 210 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 211 | return x |
| 212 | |
| 213 | |
| 214 | class ConViT(nn.Module): |