Compose with two layers
| 47 | |
| 48 | |
| 49 | class EncoderLayer(nn.Module): |
| 50 | ''' Compose with two layers ''' |
| 51 | |
| 52 | def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1): |
| 53 | super(EncoderLayer, self).__init__() |
| 54 | self.slf_attn = MultiHeadAttention( |
| 55 | n_head, d_model, d_k, d_v, dropout=dropout) |
| 56 | self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout) |
| 57 | |
| 58 | def forward(self, enc_input, non_pad_mask=None, slf_attn_mask=None): |
| 59 | enc_output, enc_slf_attn = self.slf_attn( |
| 60 | enc_input, enc_input, enc_input, mask=slf_attn_mask) |
| 61 | enc_output *= non_pad_mask |
| 62 | |
| 63 | enc_output = self.pos_ffn(enc_output) |
| 64 | enc_output *= non_pad_mask |
| 65 | |
| 66 | return enc_output, enc_slf_attn |
| 67 | |
| 68 | |
| 69 | class StarEncoderLayer(nn.Module): |