| 414 | # build a decoder layer with two multi-head attention layers and |
| 415 | # one feed-forward layer |
| 416 | class DecoderLayer(nn.Module): |
| 417 | def __init__(self, d_input, d_model, heads, dropout=0.1): |
| 418 | super().__init__() |
| 419 | self.input_linear = nn.Linear(d_input, d_model) |
| 420 | self.norm_1 = Norm(d_model) |
| 421 | self.norm_2 = Norm(d_model) |
| 422 | self.norm_3 = Norm(d_model) |
| 423 | |
| 424 | self.dropout_1 = nn.Dropout(dropout) |
| 425 | self.dropout_2 = nn.Dropout(dropout) |
| 426 | self.dropout_3 = nn.Dropout(dropout) |
| 427 | |
| 428 | self.attn_1 = MultiHeadAttention(heads, d_model, dropout=dropout) |
| 429 | self.attn_2 = MultiHeadAttention(heads, d_model, dropout=dropout) |
| 430 | self.ff = FeedForward(d_model, dropout=dropout) |
| 431 | |
| 432 | def forward(self, x, e_outputs, src_mask=None, trg_mask=None): |
| 433 | x = F.relu(self.input_linear(x)) |
| 434 | x2 = self.norm_1(x) |
| 435 | x = x + self.dropout_1(self.attn_1(x2, x2, x2, trg_mask)) |
| 436 | x2 = self.norm_2(x) |
| 437 | x = x + self.dropout_2(self.attn_2(x2, e_outputs, e_outputs, src_mask)) |
| 438 | x2 = self.norm_3(x) |
| 439 | x = x + self.dropout_3(self.ff(x2)) |
| 440 | return x |
| 441 | |
| 442 | |
| 443 | |