(self, x)
| 100 | self.proj_drop = nn.Dropout(projection_dropout) |
| 101 | |
| 102 | def forward(self, x): |
| 103 | B, N, C = x.shape |
| 104 | |
| 105 | qkv = self.qkv(x).chunk(3, dim=-1) |
| 106 | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv) |
| 107 | |
| 108 | q = q * self.scale |
| 109 | |
| 110 | attn = einsum('b h i d, b h j d -> b h i j', q, k) |
| 111 | attn = attn.softmax(dim=-1) |
| 112 | attn = self.attn_drop(attn) |
| 113 | |
| 114 | x = einsum('b h i j, b h j d -> b h i d', attn, v) |
| 115 | x = rearrange(x, 'b h n d -> b n (h d)') |
| 116 | |
| 117 | return self.proj_drop(self.proj(x)) |
| 118 | |
| 119 | |
| 120 | class TransformerEncoderLayer(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected