| 11 | class Attention(nn.Module): |
| 12 | |
| 13 | def __init__( |
| 14 | self, |
| 15 | dim, |
| 16 | num_heads=8, |
| 17 | qkv_bias=False, |
| 18 | qk_scale=None, |
| 19 | attn_drop=0.0, |
| 20 | proj_drop=0.0, |
| 21 | ): |
| 22 | super().__init__() |
| 23 | self.num_heads = num_heads |
| 24 | head_dim = dim // num_heads |
| 25 | self.scale = qk_scale or head_dim**-0.5 |
| 26 | |
| 27 | self.q = nn.Linear(dim, dim, bias=qkv_bias) |
| 28 | self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) |
| 29 | self.attn_drop = nn.Dropout(attn_drop) |
| 30 | self.proj = nn.Linear(dim, dim) |
| 31 | self.proj_drop = nn.Dropout(proj_drop) |
| 32 | |
| 33 | def forward(self, q, kv, key_mask=None): |
| 34 | N, C = kv.shape[1:] |