| 65 | return x |
| 66 | |
| 67 | class DecoderBlock(nn.Module): |
| 68 | def __init__(self, embed_dim, n_heads, dropout): |
| 69 | super(DecoderBlock, self).__init__() |
| 70 | self.self_attention = nn.MultiheadAttention(embed_dim, n_heads, dropout=dropout) |
| 71 | self.ln1 = nn.LayerNorm(embed_dim) |
| 72 | self.enc_dec_attention = nn.MultiheadAttention(embed_dim, n_heads, dropout=dropout) |
| 73 | self.ln2 = nn.LayerNorm(embed_dim) |
| 74 | self.ff = nn.Sequential( |
| 75 | nn.Linear(embed_dim, 4 * embed_dim), |
| 76 | nn.ReLU(), |
| 77 | nn.Linear(4 * embed_dim, embed_dim), |
| 78 | nn.Dropout(dropout) |
| 79 | ) |
| 80 | self.ln3 = nn.LayerNorm(embed_dim) |
| 81 | |
| 82 | def forward(self, x, enc_output, tgt_mask=None, memory_mask=None): |
| 83 | self_attn_output, _ = self.self_attention(x, x, x, attn_mask=tgt_mask) |
| 84 | x = self.ln1(x + self_attn_output) |
| 85 | enc_dec_attn_output, _ = self.enc_dec_attention(x, enc_output, enc_output, attn_mask=memory_mask) |
| 86 | x = self.ln2(x + enc_dec_attn_output) |
| 87 | ff_output = self.ff(x) |
| 88 | x = self.ln3(x + ff_output) |
| 89 | return x |
| 90 | |
| 91 | class Decoder(nn.Module): |
| 92 | def __init__(self, embed_dim, n_blocks, n_heads, dropout): |