| 66 | |
| 67 | |
| 68 | class Attention(nn.Module): |
| 69 | def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): |
| 70 | super().__init__() |
| 71 | self.num_heads = num_heads |
| 72 | head_dim = dim // num_heads |
| 73 | self.scale = qk_scale or head_dim ** -0.5 |
| 74 | |
| 75 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| 76 | self.attn_drop = nn.Dropout(attn_drop) |
| 77 | self.proj = nn.Linear(dim, dim) |
| 78 | self.proj_drop = nn.Dropout(proj_drop) |
| 79 | |
| 80 | def forward(self, x): |
| 81 | B, N, C = x.shape |
| 82 | qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) |
| 83 | q, k, v = qkv[0], qkv[1], qkv[2] |
| 84 | |
| 85 | attn = (q @ k.transpose(-2, -1)) * self.scale |
| 86 | attn = attn.softmax(dim=-1) |
| 87 | attn = self.attn_drop(attn) |
| 88 | |
| 89 | x = (attn @ v).transpose(1, 2).reshape(B, N, C) |
| 90 | x = self.proj(x) |
| 91 | x = self.proj_drop(x) |
| 92 | return x, attn |
| 93 | |
| 94 | |
| 95 | class Block(nn.Module): |