| 34 | |
| 35 | |
| 36 | class PatchMerging(nn.Module): |
| 37 | def __init__(self, patch_dim, norm_layer=nn.LayerNorm): |
| 38 | super(PatchMerging, self).__init__() |
| 39 | self.patch_dim = patch_dim |
| 40 | self.norm = norm_layer(4 * patch_dim) |
| 41 | self.reduction = nn.Linear(4 * patch_dim, 2 * patch_dim) |
| 42 | |
| 43 | def forward(self, x): |
| 44 | B, H, W, C = x.shape |
| 45 | assert H % 2 == 0 and W % 2 == 0, "Height and Width must be even." |
| 46 | |
| 47 | x0 = x[:, 0::2, 0::2, :] |
| 48 | x1 = x[:, 1::2, 0::2, :] |
| 49 | x2 = x[:, 0::2, 1::2, :] |
| 50 | x3 = x[:, 1::2, 1::2, :] |
| 51 | |
| 52 | x = torch.cat([x0, x1, x2, x3], -1) |
| 53 | x = x.view(B, -1, 4 * C) |
| 54 | x = self.norm(x) |
| 55 | x = self.reduction(x) |
| 56 | x = x.view(B, H // 2, W // 2, 2 * C) |
| 57 | return x |
| 58 | |
| 59 | |
| 60 | class FeatureExtractor(nn.Module): |