| 595 | # Token merging and splitting |
| 596 | |
| 597 | class TokenMerge(nn.Module): |
| 598 | def __init__(self, in_features, out_features, patch_size=(2, 2)): |
| 599 | super().__init__() |
| 600 | self.h = patch_size[0] |
| 601 | self.w = patch_size[1] |
| 602 | self.proj = apply_wd(Linear(in_features * self.h * self.w, out_features, bias=False)) |
| 603 | |
| 604 | def forward(self, x): |
| 605 | x = rearrange(x, "... (h nh) (w nw) e -> ... h w (nh nw e)", nh=self.h, nw=self.w) |
| 606 | x=self.proj(x) |
| 607 | return x |
| 608 | |
| 609 | class TokenSplitWithoutSkip(nn.Module): |
| 610 | def __init__(self, in_features, out_features, patch_size=(2, 2)): |