| 5 | from typing import Optional |
| 6 | |
| 7 | class MultiHeadAttention(nn.Module): |
| 8 | def __init__(self, embed_dim, num_heads, dropout=0.1, bias=True): |
| 9 | super().__init__() |
| 10 | assert embed_dim % num_heads == 0 |
| 11 | |
| 12 | self.embed_dim = embed_dim |
| 13 | self.num_heads = num_heads |
| 14 | self.head_dim = embed_dim // num_heads |
| 15 | self.scale = self.head_dim ** -0.5 |
| 16 | |
| 17 | self.qkv = nn.Linear(embed_dim, embed_dim * 3, bias=bias) |
| 18 | self.proj = nn.Linear(embed_dim, embed_dim, bias=bias) |
| 19 | self.dropout = nn.Dropout(dropout) |
| 20 | |
| 21 | def forward(self, x, mask=None): |
| 22 | B, N, C = x.shape |
| 23 | |
| 24 | qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) |
| 25 | q, k, v = qkv[0], qkv[1], qkv[2] |
| 26 | |
| 27 | attn = (q @ k.transpose(-2, -1)) * self.scale |
| 28 | |
| 29 | if mask is not None: |
| 30 | attn = attn.masked_fill(mask == 0, float('-inf')) |
| 31 | |
| 32 | attn = F.softmax(attn, dim=-1) |
| 33 | attn = self.dropout(attn) |
| 34 | |
| 35 | x = (attn @ v).transpose(1, 2).reshape(B, N, C) |
| 36 | x = self.proj(x) |
| 37 | x = self.dropout(x) |
| 38 | |
| 39 | return x |
| 40 | |
| 41 | class CrossAttention(nn.Module): |
| 42 | def __init__(self, embed_dim, num_heads, dropout=0.1): |