Args: x: (B, H, W, C) window_size (int): window size Returns: windows: (num_windows*B, window_size, window_size, C)
(x, window_size)
| 31 | |
| 32 | |
| 33 | def window_partition(x, window_size): |
| 34 | """ |
| 35 | Args: |
| 36 | x: (B, H, W, C) |
| 37 | window_size (int): window size |
| 38 | |
| 39 | Returns: |
| 40 | windows: (num_windows*B, window_size, window_size, C) |
| 41 | """ |
| 42 | B, H, W, C = x.shape |
| 43 | x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) |
| 44 | windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) |
| 45 | return windows |
| 46 | |
| 47 | |
| 48 | def window_reverse(windows, window_size, H, W): |