| 62 | |
| 63 | |
| 64 | def window_partition(x, window_size): |
| 65 | B, H, W, C = x.shape |
| 66 | pad_h = (window_size - H % window_size) % window_size |
| 67 | pad_w = (window_size - W % window_size) % window_size |
| 68 | if pad_h > 0 or pad_w > 0: |
| 69 | x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) |
| 70 | Hp, Wp = H + pad_h, W + pad_w |
| 71 | x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) |
| 72 | windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C) |
| 73 | return windows, (Hp, Wp) |
| 74 | |
| 75 | |
| 76 | def window_unpartition(windows, window_size, pad_hw, hw): |