| 35 | return x |
| 36 | |
| 37 | class EncoderBlock(nn.Module): |
| 38 | def __init__(self, embed_dim, n_heads, dropout): |
| 39 | super(EncoderBlock, self).__init__() |
| 40 | self.attention = nn.MultiheadAttention(embed_dim, n_heads, dropout=dropout) |
| 41 | self.ln1 = nn.LayerNorm(embed_dim) |
| 42 | self.ff = nn.Sequential( |
| 43 | nn.Linear(embed_dim, 4 * embed_dim), |
| 44 | nn.ReLU(), |
| 45 | nn.Linear(4 * embed_dim, embed_dim), |
| 46 | nn.Dropout(dropout) |
| 47 | ) |
| 48 | self.ln2 = nn.LayerNorm(embed_dim) |
| 49 | |
| 50 | def forward(self, x, src_mask=None): |
| 51 | attn_output, _ = self.attention(x, x, x, attn_mask=src_mask) |
| 52 | x = self.ln1(x + attn_output) |
| 53 | ff_output = self.ff(x) |
| 54 | x = self.ln2(x + ff_output) |
| 55 | return x |
| 56 | |
| 57 | class Encoder(nn.Module): |
| 58 | def __init__(self, embed_dim, n_blocks, n_heads, dropout): |