| 19 | class LocalAttention(nn.Module): |
| 20 | |
| 21 | def __init__(self, dim, heads, window_size, attn_drop, proj_drop): |
| 22 | |
| 23 | super().__init__() |
| 24 | |
| 25 | window_size = to_2tuple(window_size) |
| 26 | |
| 27 | self.proj_qkv = nn.Linear(dim, 3 * dim) |
| 28 | self.heads = heads |
| 29 | assert dim % heads == 0 |
| 30 | head_dim = dim // heads |
| 31 | self.scale = head_dim ** -0.5 |
| 32 | self.proj_out = nn.Linear(dim, dim) |
| 33 | self.window_size = window_size |
| 34 | self.proj_drop = nn.Dropout(proj_drop, inplace=True) |
| 35 | self.attn_drop = nn.Dropout(attn_drop, inplace=True) |
| 36 | |
| 37 | Wh, Ww = self.window_size |
| 38 | self.relative_position_bias_table = nn.Parameter( |
| 39 | torch.zeros((2 * Wh - 1) * (2 * Ww - 1), heads) |
| 40 | ) |
| 41 | trunc_normal_(self.relative_position_bias_table, std=0.01) |
| 42 | |
| 43 | coords_h = torch.arange(self.window_size[0]) |
| 44 | coords_w = torch.arange(self.window_size[1]) |
| 45 | coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww |
| 46 | coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww |
| 47 | relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww |
| 48 | relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 |
| 49 | relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 |
| 50 | relative_coords[:, :, 1] += self.window_size[1] - 1 |
| 51 | relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 |
| 52 | relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww |
| 53 | self.register_buffer("relative_position_index", relative_position_index) |
| 54 | |
| 55 | def forward(self, x, mask=None): |
| 56 | |