| 220 | |
| 221 | |
| 222 | class TransformerDecoderLayer(layer.Layer): |
| 223 | def __init__(self, d_model=512, n_head=8, dim_feedforward=2048): |
| 224 | super(TransformerDecoderLayer, self).__init__() |
| 225 | |
| 226 | self.d_model = d_model |
| 227 | self.n_head = n_head |
| 228 | self.dim_feedforward = dim_feedforward |
| 229 | |
| 230 | self.dec_self_attn = MultiHeadAttention(d_model=d_model, n_head=n_head) |
| 231 | self.dec_enc_attn = MultiHeadAttention(d_model=d_model, n_head=n_head) |
| 232 | self.pos_ffn = PoswiseFeedForwardNet(d_model=d_model, dim_feedforward=dim_feedforward) |
| 233 | |
| 234 | def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask): |
| 235 | """ |
| 236 | Args: |
| 237 | dec_inputs: [batch_size, tgt_len, d_model] |
| 238 | enc_outputs: [batch_size, src_len, d_model] |
| 239 | dec_self_attn_mask: [batch_size, tgt_len, tgt_len] |
| 240 | dec_enc_attn_mask: [batch_size, tgt_len, src_len] |
| 241 | """ |
| 242 | |
| 243 | # dec_outputs: [batch_size, tgt_len, d_model] |
| 244 | # dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len] |
| 245 | dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask) |
| 246 | |
| 247 | # dec_outputs: [batch_size, tgt_len, d_model] |
| 248 | # dec_self_attn: [batch_size, n_heads, tgt_len, src_len] |
| 249 | dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask) |
| 250 | # [batch_size, tgt_len, d_model] |
| 251 | dec_outputs = self.pos_ffn(dec_outputs) |
| 252 | return dec_outputs, dec_self_attn, dec_enc_attn |
| 253 | |
| 254 | |
| 255 | class TransformerEncoder(layer.Layer): |