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)
| 56 | |
| 57 | |
| 58 | def window_reverse(windows, window_size, H, W): |
| 59 | """ |
| 60 | Args: |
| 61 | windows: (num_windows*B, window_size, window_size, C) |
| 62 | window_size (int): Window size |
| 63 | H (int): Height of image |
| 64 | W (int): Width of image |
| 65 | Returns: |
| 66 | x: (B, H, W, C) |
| 67 | """ |
| 68 | B = int(windows.shape[0] / (H * W / window_size / window_size)) |
| 69 | x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) |
| 70 | x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) |
| 71 | return x |
| 72 | |
| 73 | |
| 74 | class WindowAttention(nn.Module): |