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