| 88 | # modules |
| 89 | |
| 90 | class Attention(nn.Module): |
| 91 | def __init__(self, dim, num_heads=8, attention_dropout=0.1, projection_dropout=0.1): |
| 92 | super().__init__() |
| 93 | self.heads = num_heads |
| 94 | head_dim = dim // self.heads |
| 95 | self.scale = head_dim ** -0.5 |
| 96 | |
| 97 | self.qkv = nn.Linear(dim, dim * 3, bias=False) |
| 98 | self.attn_drop = nn.Dropout(attention_dropout) |
| 99 | self.proj = nn.Linear(dim, dim) |
| 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): |