(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
window_size=None, window=False)
| 249 | |
| 250 | class Block(nn.Module): |
| 251 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., |
| 252 | drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm, |
| 253 | window_size=None, window=False): |
| 254 | super().__init__() |
| 255 | self.norm1 = norm_layer(dim) |
| 256 | if not window: |
| 257 | self.attn = Attention( |
| 258 | dim, num_heads=num_heads, qkv_bias=qkv_bias, |
| 259 | attn_drop=attn_drop, proj_drop=drop, window_size=window_size) |
| 260 | else: |
| 261 | self.attn = WindowAttention( |
| 262 | dim, num_heads=num_heads, qkv_bias=qkv_bias, |
| 263 | attn_drop=attn_drop, proj_drop=drop, window_size=window_size) |
| 264 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
| 265 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 266 | self.norm2 = norm_layer(dim) |
| 267 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 268 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 269 | |
| 270 | if init_values is not None: |
| 271 | self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) |
| 272 | self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) |
| 273 | else: |
| 274 | self.gamma_1, self.gamma_2 = None, None |
| 275 | |
| 276 | def forward(self, x, H, W): |
| 277 | if self.gamma_1 is None: |
nothing calls this directly
no test coverage detected