Unpatchify a tensor. Args: x (torch.Tensor): (N, C, *spatial) tensor patch_size (int): Patch size
(x: torch.Tensor, patch_size: int)
| 32 | |
| 33 | |
| 34 | def unpatchify(x: torch.Tensor, patch_size: int): |
| 35 | """ |
| 36 | Unpatchify a tensor. |
| 37 | |
| 38 | Args: |
| 39 | x (torch.Tensor): (N, C, *spatial) tensor |
| 40 | patch_size (int): Patch size |
| 41 | """ |
| 42 | DIM = x.dim() - 2 |
| 43 | assert x.shape[1] % (patch_size ** DIM) == 0, f"Second dimension of input tensor must be divisible by patch size to unpatchify, got {x.shape[1]} and {patch_size ** DIM}" |
| 44 | |
| 45 | x = x.reshape(x.shape[0], x.shape[1] // (patch_size ** DIM), *([patch_size] * DIM), *(x.shape[-DIM:])) |
| 46 | x = x.permute(0, 1, *(sum([[2 + DIM + i, 2 + i] for i in range(DIM)], []))) |
| 47 | x = x.reshape(x.shape[0], x.shape[1], *[x.shape[2 + 2 * i] * patch_size for i in range(DIM)]) |
| 48 | return x |