| 137 | |
| 138 | |
| 139 | class SwinBlock(nn.Module): |
| 140 | def __init__(self, dim, heads, head_dim, mlp_dim, shifted, window_size, relative_pos_embedding): |
| 141 | super().__init__() |
| 142 | self.attention_block = Residual(PreNorm(dim, WindowAttention(dim=dim, |
| 143 | heads=heads, |
| 144 | head_dim=head_dim, |
| 145 | shifted=shifted, |
| 146 | window_size=window_size, |
| 147 | relative_pos_embedding=relative_pos_embedding))) |
| 148 | self.mlp_block = Residual(PreNorm(dim, FeedForward(dim=dim, hidden_dim=mlp_dim))) |
| 149 | |
| 150 | def forward(self, x): |
| 151 | x = self.attention_block(x) |
| 152 | x = self.mlp_block(x) |
| 153 | return x |
| 154 | |
| 155 | |
| 156 | class PatchMerging(nn.Module): |