Args: windows: (num_windows*B, window_size, window_size, C) window_size (int): Window size H (int): Height of image W (int): Width of image Returns: x: (B, H, W, C)
(windows, window_size, H, W)
| 61 | |
| 62 | |
| 63 | def window_reverse(windows, window_size, H, W): |
| 64 | """ |
| 65 | Args: |
| 66 | windows: (num_windows*B, window_size, window_size, C) |
| 67 | window_size (int): Window size |
| 68 | H (int): Height of image |
| 69 | W (int): Width of image |
| 70 | Returns: |
| 71 | x: (B, H, W, C) |
| 72 | """ |
| 73 | B = int(windows.shape[0] / (H * W / window_size / window_size)) |
| 74 | x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) |
| 75 | x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) |
| 76 | return x |
| 77 | |
| 78 | |
| 79 | class WindowAttention(nn.Module): |