将TensorFlow 1 contrib代码迁移为TF2 Keras兼容版本的问题求助

我正在将代码从TF1迁移到TF2，目前已经修复了绝大多数适配TF2运行的问题，但在迁移到可兼容tfa.seq2seq.LuongAttention和tfa.seq2seq.AttentionWrapper的TF2版本时遇到阻碍。我已将contrib相关接口替换为v2版本，但无法定位运行异常的原因。
原有问题代码如下：

def _single_cell(num_units, keep_prob, device_str=None):
    single_cell = tf.compat.v1.nn.rnn_cell.GRUCell(num_units)
    if keep_prob < 1.0:
        single_cell = tf.contrib.rnn.DropoutWrapper(cell=single_cell, input_keep_prob=keep_prob)
    # Device Wrapper
    if device_str:
        single_cell = tf.contrib.rnn.DeviceWrapper(single_cell, device_str)
    return single_cell


def create_rnn_cell(num_units, num_layers, keep_prob):
    """Create multi-layer RNN cell."""
    cell_list = []
    for i in range(num_layers):
        single_cell = _single_cell(num_units=num_units, keep_prob=keep_prob)
        cell_list.append(single_cell)
    if len(cell_list) == 1:  # Single layer.
        return cell_list[0]
    else:  # Multi layers
        return tf.compat.v1.nn.rnn_cell.MultiRNNCell(cell_list)

cell = create_rnn_cell(
            num_units=hparams.num_units,
            num_layers=hparams.num_layers,
            keep_prob=hparams.keep_prob)

encoder_outputs, encoder_state = tf.compat.v1.nn.dynamic_rnn(
                cell,
                encoder_emb_inp,
                dtype=dtype,
                sequence_length=self.batch_input.source_sequence_length,
                time_major=self.time_major)

我参考了TFA官方seq2seq模块的文档，已经完成了大部分代码的迁移工作，仅剩下上述两个函数无法完成适配。

修改适配说明 #

核心调整点 #

• 所有tf.compat.v1.nn.rnn_cell、tf.contrib.rnn下的TF1风格RNN接口全部替换为TF2原生tf.keras.layers下的标准RNN实现，TFA的seq2seq相关组件仅兼容Keras标准RNN Cell
• Dropout功能直接通过Keras GRUCell的内置参数实现，不需要单独用DropoutWrapper
• 设备绑定功能可以通过tf.device上下文管理器实现，无需使用contrib的DeviceWrapper
• TF1的dynamic_rnn接口替换为tf.keras.layers.RNN层的调用逻辑，序列长度通过mask参数传入

适配后的完整代码 #

import tensorflow as tf
import tensorflow_addons as tfa

def _single_cell(num_units, keep_prob, device_str=None):
    # 替换为TF2原生Keras GRUCell，直接配置dropout参数
    single_cell = tf.keras.layers.GRUCell(
        units=num_units,
        dropout=1 - keep_prob if keep_prob < 1.0 else 0.0
    )
    # 设备绑定
    if device_str:
        with tf.device(device_str):
            single_cell.build(input_shape=(None, num_units))
    return single_cell


def create_rnn_cell(num_units, num_layers, keep_prob):
    """Create multi-layer RNN cell."""
    cell_list = []
    for i in range(num_layers):
        single_cell = _single_cell(num_units=num_units, keep_prob=keep_prob)
        cell_list.append(single_cell)
    if len(cell_list) == 1:  # 单层级直接返回
        return cell_list[0]
    else:  # 多层级用StackedRNNCells替代TF1的MultiRNNCell
        return tf.keras.layers.StackedRNNCells(cell_list)

# 实例化RNN Cell
cell = create_rnn_cell(
            num_units=hparams.num_units,
            num_layers=hparams.num_layers,
            keep_prob=hparams.keep_prob)

# 替换dynamic_rnn为Keras RNN层调用
rnn_layer = tf.keras.layers.RNN(
    cell,
    return_sequences=True,
    return_state=True,
    time_major=self.time_major
)
encoder_outputs, encoder_state = rnn_layer(
    encoder_emb_inp,
    mask=tf.sequence_mask(self.batch_input.source_sequence_length),
    dtype=dtype
)

对接TFA注意力模块示例 #

修改后的Cell可以直接适配TFA的注意力组件：

# 实例化Luong注意力机制
attention_mechanism = tfa.seq2seq.LuongAttention(
    units=hparams.num_units,
    memory=encoder_outputs,
    memory_sequence_length=self.batch_input.source_sequence_length
)

# 用AttentionWrapper包裹RNN Cell即可完成对接
attention_cell = tfa.seq2seq.AttentionWrapper(
    cell=cell,
    attention_mechanism=attention_mechanism,
    attention_layer_size=hparams.num_units
)

内容的提问来源于stack exchange，提问作者include conio.h