Patchify a tensor. Args: x (torch.Tensor): (N, C, *spatial) tensor patch_size (int): Patch size
(x: torch.Tensor, patch_size: int)
| 14 | |
| 15 | |
| 16 | def patchify(x: torch.Tensor, patch_size: int): |
| 17 | """ |
| 18 | Patchify a tensor. |
| 19 | |
| 20 | Args: |
| 21 | x (torch.Tensor): (N, C, *spatial) tensor |
| 22 | patch_size (int): Patch size |
| 23 | """ |
| 24 | DIM = x.dim() - 2 |
| 25 | for d in range(2, DIM + 2): |
| 26 | assert x.shape[d] % patch_size == 0, f"Dimension {d} of input tensor must be divisible by patch size, got {x.shape[d]} and {patch_size}" |
| 27 | |
| 28 | x = x.reshape(*x.shape[:2], *sum([[x.shape[d] // patch_size, patch_size] for d in range(2, DIM + 2)], [])) |
| 29 | x = x.permute(0, 1, *([2 * i + 3 for i in range(DIM)] + [2 * i + 2 for i in range(DIM)])) |
| 30 | x = x.reshape(x.shape[0], x.shape[1] * (patch_size ** DIM), *(x.shape[-DIM:])) |
| 31 | return x |
| 32 | |
| 33 | |
| 34 | def unpatchify(x: torch.Tensor, patch_size: int): |