Args: x: (B, H, W, C) window_size (int): window size Returns: windows: (num_windows*B, window_size, window_size, C)
(x, window_size)
| 47 | |
| 48 | |
| 49 | def window_partition(x, window_size): |
| 50 | """ |
| 51 | Args: |
| 52 | x: (B, H, W, C) |
| 53 | window_size (int): window size |
| 54 | Returns: |
| 55 | windows: (num_windows*B, window_size, window_size, C) |
| 56 | """ |
| 57 | B, H, W, C = x.shape |
| 58 | x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) |
| 59 | windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) |
| 60 | return windows |
| 61 | |
| 62 | |
| 63 | def window_reverse(windows, window_size, H, W): |