(self, x)
| 206 | self.proj_drop = nn.Dropout(proj_drop) |
| 207 | |
| 208 | def forward(self, x): |
| 209 | B, N, C = x.shape |
| 210 | qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) |
| 211 | q, k, v = qkv[0], qkv[1], qkv[2] |
| 212 | |
| 213 | attn = (q @ k.transpose(-2, -1)) * self.scale |
| 214 | attn = attn.softmax(dim=-1) |
| 215 | attn = self.attn_drop(attn) |
| 216 | |
| 217 | x = (attn @ v).transpose(1, 2).reshape(B, N, C) |
| 218 | x = self.proj(x) |
| 219 | x = self.proj_drop(x) |
| 220 | return x, attn |
| 221 | |
| 222 | |
| 223 | class Block(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected