| 169 | |
| 170 | |
| 171 | class StageModule(nn.Module): |
| 172 | def __init__(self, in_channels, hidden_dimension, layers, downscaling_factor, num_heads, head_dim, window_size, |
| 173 | relative_pos_embedding): |
| 174 | super().__init__() |
| 175 | assert layers % 2 == 0, 'Stage layers need to be divisible by 2 for regular and shifted block.' |
| 176 | |
| 177 | self.patch_partition = PatchMerging(in_channels=in_channels, out_channels=hidden_dimension, |
| 178 | downscaling_factor=downscaling_factor) |
| 179 | |
| 180 | self.layers = nn.ModuleList([]) |
| 181 | for _ in range(layers // 2): |
| 182 | self.layers.append(nn.ModuleList([ |
| 183 | SwinBlock(dim=hidden_dimension, heads=num_heads, head_dim=head_dim, mlp_dim=hidden_dimension * 4, |
| 184 | shifted=False, window_size=window_size, relative_pos_embedding=relative_pos_embedding), |
| 185 | SwinBlock(dim=hidden_dimension, heads=num_heads, head_dim=head_dim, mlp_dim=hidden_dimension * 4, |
| 186 | shifted=True, window_size=window_size, relative_pos_embedding=relative_pos_embedding), |
| 187 | ])) |
| 188 | |
| 189 | def forward(self, x): |
| 190 | x = self.patch_partition(x) |
| 191 | for regular_block, shifted_block in self.layers: |
| 192 | x = regular_block(x) |
| 193 | x = shifted_block(x) |
| 194 | return x.permute(0, 3, 1, 2) |
| 195 | |
| 196 | |
| 197 | class SwinTransformer(nn.Module): |