Method forward

examples/singa_peft/examples/model/trans.py:162–202

Args: dec_inputs: [batch_size, tgt_len] enc_inputs: [batch_size, src_len] enc_outputs: [batch_size, src_len, d_model]

(self, dec_inputs, enc_inputs, enc_outputs)

            self.layers.append(TransformerDecoderLayer(d_model=d_model, n_head=n_head, dim_feedforward=dim_feedforward))

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        """
        Args:
            dec_inputs: [batch_size, tgt_len]
            enc_inputs: [batch_size, src_len]
            enc_outputs: [batch_size, src_len, d_model]

        """

        # [batch_size, tgt_len, d_model]
        tgt_word_emb = self.target_emb(dec_inputs)
        self.target_pos_emb.initialize(dec_inputs)
        self.target_pos_emb.from_pretrained(W=TransformerDecoder._get_sinusoid_encoding_table(self.tgt_n_token, self.d_model),
                                            freeze=True)
        # [batch_size, tgt_len, d_model]
        tgt_pos_emb = self.target_pos_emb(dec_inputs)
        # [batch_size, tgt_len, d_model]
        dec_outputs = autograd.add(tgt_word_emb, tgt_pos_emb)

        # dec_self_attn_pad_mask  [batch_size, tgt_len, tgt_len]
        dec_self_attn_pad_mask = TransformerDecoder._get_attn_pad_mask(dec_inputs, dec_inputs)
        # [batch_size, tgt_len, tgt_len]
        dec_self_attn_subsequent_mask = TransformerDecoder._get_attn_subsequence_mask(dec_inputs)

        # dec_self_attn_mask [batch_size, tgt_len, tgt_len]
        dec_self_attn_mask = tensor.gt((dec_self_attn_pad_mask + dec_self_attn_subsequent_mask), 0)

        # dec_enc_attn_mask [batch_size, tgt_len, src_len]
        dec_enc_attn_mask = TransformerDecoder._get_attn_pad_mask(dec_inputs, enc_inputs)

        dec_self_attns, dec_enc_attns = [], []

        for layer in self.layers:
            # dec_outputs: [batch_size, tgt_len, d_model],
            # dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len],
            # dec_enc_attn: [batch_size, h_heads, tgt_len,src_len]
            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask,
                                                             dec_enc_attn_mask)
            dec_self_attns.append(dec_self_attn)
            dec_enc_attns.append(dec_enc_attn)
        return dec_outputs, dec_self_attns, dec_enc_attns

    @staticmethod
    def _get_attn_pad_mask(seq_q, seq_k):