| 17 | from timm.models.layers import to_2tuple, trunc_normal_ |
| 18 | |
| 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 | |
| 57 | B, C, H, W = x.size() |
| 58 | r1, r2 = H // self.window_size[0], W // self.window_size[1] |
| 59 | |
| 60 | x_total = einops.rearrange(x, 'b c (r1 h1) (r2 w1) -> b (r1 r2) (h1 w1) c', h1=self.window_size[0], w1=self.window_size[1]) # B x Nr x Ws x C |
| 61 | |
| 62 | x_total = einops.rearrange(x_total, 'b m n c -> (b m) n c') |
| 63 | |
| 64 | qkv = self.proj_qkv(x_total) # B' x N x 3C |
| 65 | q, k, v = torch.chunk(qkv, 3, dim=2) |
| 66 | |
| 67 | q = q * self.scale |
| 68 | q, k, v = [einops.rearrange(t, 'b n (h c1) -> b h n c1', h=self.heads) for t in [q, k, v]] |
| 69 | attn = torch.einsum('b h m c, b h n c -> b h m n', q, k) |
| 70 | |
| 71 | relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( |
| 72 | self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH |
| 73 | relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww |
| 74 | attn_bias = relative_position_bias |
| 75 | attn = attn + attn_bias.unsqueeze(0) |
| 76 | |