Swin Transformer Block. Args: dim (int): Number of input channels. num_heads (int): Number of attention heads. window_size (int): Window size. shift_size (int): Shift size for SW-MSA. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. qkv
| 178 | |
| 179 | |
| 180 | class SwinTransformerBlock(nn.Module): |
| 181 | """Swin Transformer Block. |
| 182 | Args: |
| 183 | dim (int): Number of input channels. |
| 184 | num_heads (int): Number of attention heads. |
| 185 | window_size (int): Window size. |
| 186 | shift_size (int): Shift size for SW-MSA. |
| 187 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. |
| 188 | qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True |
| 189 | qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. |
| 190 | drop (float, optional): Dropout rate. Default: 0.0 |
| 191 | attn_drop (float, optional): Attention dropout rate. Default: 0.0 |
| 192 | drop_path (float, optional): Stochastic depth rate. Default: 0.0 |
| 193 | act_layer (nn.Module, optional): Activation layer. Default: nn.GELU |
| 194 | norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm |
| 195 | """ |
| 196 | |
| 197 | def __init__( |
| 198 | self, |
| 199 | dim, |
| 200 | num_heads, |
| 201 | window_size=7, |
| 202 | shift_size=0, |
| 203 | mlp_ratio=4.0, |
| 204 | qkv_bias=True, |
| 205 | qk_scale=None, |
| 206 | drop=0.0, |
| 207 | attn_drop=0.0, |
| 208 | drop_path=0.0, |
| 209 | act_layer=nn.GELU, |
| 210 | norm_layer=nn.LayerNorm, |
| 211 | ): |
| 212 | super().__init__() |
| 213 | self.dim = dim |
| 214 | self.num_heads = num_heads |
| 215 | self.window_size = window_size |
| 216 | self.shift_size = shift_size |
| 217 | self.mlp_ratio = mlp_ratio |
| 218 | assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" |
| 219 | |
| 220 | self.norm1 = norm_layer(dim) |
| 221 | self.attn = WindowAttention( |
| 222 | dim, |
| 223 | window_size=to_2tuple(self.window_size), |
| 224 | num_heads=num_heads, |
| 225 | qkv_bias=qkv_bias, |
| 226 | qk_scale=qk_scale, |
| 227 | attn_drop=attn_drop, |
| 228 | proj_drop=drop, |
| 229 | ) |
| 230 | |
| 231 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 232 | self.norm2 = norm_layer(dim) |
| 233 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 234 | self.mlp = Mlp( |
| 235 | in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop |
| 236 | ) |
| 237 |