| 62 | |
| 63 | class Attention(nn.Module): |
| 64 | def __init__( |
| 65 | self, dim, num_heads=8, qkv_bias=False, attn_drop=0., |
| 66 | proj_drop=0., window_size=None): |
| 67 | super().__init__() |
| 68 | self.num_heads = num_heads |
| 69 | head_dim = dim // num_heads |
| 70 | self.scale = head_dim ** -0.5 |
| 71 | |
| 72 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| 73 | self.window_size = window_size |
| 74 | q_size = window_size[0] |
| 75 | kv_size = q_size |
| 76 | rel_sp_dim = 2 * q_size - 1 |
| 77 | self.rel_pos_h = nn.Parameter(torch.zeros(rel_sp_dim, head_dim)) |
| 78 | self.rel_pos_w = nn.Parameter(torch.zeros(rel_sp_dim, head_dim)) |
| 79 | |
| 80 | self.attn_drop = nn.Dropout(attn_drop) |
| 81 | self.proj = nn.Linear(dim, dim) |
| 82 | self.proj_drop = nn.Dropout(proj_drop) |
| 83 | |
| 84 | def forward(self, x, H, W, rel_pos_bias=None): |
| 85 | B, N, C = x.shape |