x: [B, L, C].
(self, x)
| 361 | nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) |
| 362 | |
| 363 | def forward(self, x): |
| 364 | """ |
| 365 | x: [B, L, C]. |
| 366 | """ |
| 367 | b, s, c, n, d = *x.size(), self.num_heads, self.head_dim |
| 368 | |
| 369 | # compute query, key, value |
| 370 | q = self.to_q(self.cls_embedding).view(1, 1, n*d).expand(b, -1, -1) |
| 371 | k, v = self.to_kv(x).chunk(2, dim=-1) |
| 372 | |
| 373 | # compute attention |
| 374 | x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) |
| 375 | x = x.reshape(b, 1, c) |
| 376 | |
| 377 | # output |
| 378 | x = self.proj(x) |
| 379 | x = F.dropout(x, self.proj_dropout, self.training) |
| 380 | |
| 381 | # mlp |
| 382 | x = x + self.mlp(self.norm(x)) |
| 383 | return x[:, 0] |
| 384 | |
| 385 | |
| 386 | class VisionTransformer(nn.Module): |
nothing calls this directly
no test coverage detected