| 186 | """ |
| 187 | |
| 188 | def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.): |
| 189 | super().__init__() |
| 190 | self.dim = dim |
| 191 | self.window_size = window_size # Wh, Ww |
| 192 | self.num_heads = num_heads |
| 193 | head_dim = dim // num_heads |
| 194 | self.scale = head_dim ** -0.5 |
| 195 | |
| 196 | q_size = window_size[0] |
| 197 | kv_size = window_size[1] |
| 198 | rel_sp_dim = 2 * q_size - 1 |
| 199 | self.rel_pos_h = nn.Parameter(torch.zeros(rel_sp_dim, head_dim)) |
| 200 | self.rel_pos_w = nn.Parameter(torch.zeros(rel_sp_dim, head_dim)) |
| 201 | |
| 202 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| 203 | self.attn_drop = nn.Dropout(attn_drop) |
| 204 | self.proj = nn.Linear(dim, dim) |
| 205 | self.proj_drop = nn.Dropout(proj_drop) |
| 206 | |
| 207 | def forward(self, x, H, W): |
| 208 | """ Forward function. |