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