python - 为堆叠的双向编码器-解码器 seq 添加波束解码器到 seq 模型

Question

我希望在 TF-keras 中实现的 NMT 的 seq2seq 模型的编码器-解码器堆叠 LSTM 层的末尾进行波束搜索。我在用

在解码器类

self.decoder_rnn_cell = tf.keras.layers.LSTMCell(self.dec_units)

def build_rnn_cell(self, batch_sz):
    rnn_cell = tfa.seq2seq.AttentionWrapper(self.decoder_rnn_cell,
                                  self.attention_mechanism, attention_layer_size=self.dec_units)
    return rnn_cell

self.rnn_cell = self.build_rnn_cell(batch_sz)

#in beam_search_evaluate

decoder_instance = tfa.seq2seq.BeamSearchDecoder(cell=decoder.rnn_cell, beam_width=beam_width, output_layer=decoder.fc, maximum_iterations=max_target_len)

我需要将解码器 RNN 单元更改为堆叠层解码器。我不确定如何将此decoder.rnn_cell 转换为堆叠层LSTM。我应该使用 keras.stackedRNNCell 层吗？任何在解码器中具有堆叠 RNN 层的 TFa 波束搜索解码器的代码片段（使用 tfa.seq2seq.AttentionWrapper）都会非常有帮助。