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. q
| 144 | |
| 145 | |
| 146 | class SwinTransformerBlock(nn.Module): |
| 147 | """ Swin Transformer Block. |
| 148 | |
| 149 | Args: |
| 150 | dim (int): Number of input channels. |
| 151 | num_heads (int): Number of attention heads. |
| 152 | window_size (int): Window size. |
| 153 | shift_size (int): Shift size for SW-MSA. |
| 154 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. |
| 155 | qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True |
| 156 | qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. |
| 157 | drop (float, optional): Dropout rate. Default: 0.0 |
| 158 | attn_drop (float, optional): Attention dropout rate. Default: 0.0 |
| 159 | drop_path (float, optional): Stochastic depth rate. Default: 0.0 |
| 160 | act_layer (nn.Module, optional): Activation layer. Default: nn.GELU |
| 161 | norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm |
| 162 | """ |
| 163 | |
| 164 | def __init__(self, dim, num_heads, window_size=7, shift_size=0, |
| 165 | mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., |
| 166 | act_layer=nn.GELU, norm_layer=nn.LayerNorm): |
| 167 | super().__init__() |
| 168 | self.dim = dim |
| 169 | self.num_heads = num_heads |
| 170 | self.window_size = window_size |
| 171 | self.shift_size = shift_size |
| 172 | self.mlp_ratio = mlp_ratio |
| 173 | assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" |
| 174 | |
| 175 | self.norm1 = norm_layer(dim) |
| 176 | self.attn = WindowAttention( |
| 177 | dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, |
| 178 | qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) |
| 179 | |
| 180 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 181 | self.norm2 = norm_layer(dim) |
| 182 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 183 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 184 | |
| 185 | self.H = None |
| 186 | self.W = None |
| 187 | |
| 188 | def forward(self, x, mask_matrix): |
| 189 | """ Forward function. |
| 190 | |
| 191 | Args: |
| 192 | x: Input feature, tensor size (B, H*W, C). |
| 193 | H, W: Spatial resolution of the input feature. |
| 194 | mask_matrix: Attention mask for cyclic shift. |
| 195 | """ |
| 196 | B, L, C = x.shape |
| 197 | H, W = self.H, self.W |
| 198 | assert L == H * W, "input feature has wrong size" |
| 199 | |
| 200 | shortcut = x |
| 201 | x = self.norm1(x) |
| 202 | x = x.view(B, H, W, C) |
| 203 |