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