将tf.compat.v1.nn.dynamic_rnn迁移至TensorFlow 2.0的tf.keras.layers.RNN

问题背景 #

旧代码基于TensorFlow 1.x的tf.compat.v1.nn.dynamic_rnn实现RNN功能，在TensorFlow 2.x环境下无法运行，需要迁移到tf.keras.layers.RNN。已知堆叠cell定义为drop_multi_cell = tf.keras.layers.StackedRNNCells(drop_lstm_cells)，但测试实现存在错误。

原TF1.x代码：

all_lstm_outputs, self.state = tf.compat.v1.nn.dynamic_rnn(
        drop_multi_cell, all_inputs, initial_state=tuple(initial_state),
        time_major = True, dtype=tf.float32)

测试代码的问题 #

测试代码中tf.keras.layers.RNN([drop_multi_cell])的写法错误——StackedRNNCells本身就是多个RNN cell的堆叠封装，属于单个RNNCell实例，不需要用列表包裹传入。

正确实现代码 #

# 创建RNN层：直接传入StackedRNNCells实例，无需列表包裹
rnn_layer = tf.keras.layers.RNN(drop_multi_cell, time_major=True, dtype=tf.float32)
# 调用层获取输出与最终状态
all_lstm_outputs, self.state = rnn_layer(all_inputs, initial_state=tuple(initial_state))

参数对应说明 #

• cell参数：原dynamic_rnn接受堆叠cell（如MultiRNNCell），Keras RNN层直接传入StackedRNNCells实例即可，无需额外封装列表
• time_major：保持与原代码一致的True，确保输入维度顺序为[时间步, 批量大小, 特征数]
• initial_state：原代码的tuple(initial_state)是各LSTM层初始状态（每个状态为(c, h)元组）组成的元组，与Keras RNN层要求的格式完全匹配，直接传入即可
• dtype：在创建RNN层时指定dtype=tf.float32，对应原代码的dtype参数，确保输出和状态的数据类型一致

核心API差异说明（官方文档核心内容翻译） #

tf.compat.v1.nn.dynamic_rnn #

动态构建RNN以处理可变长度序列，支持time_major输入格式，可接受单个RNNCell或堆叠型cell（如MultiRNNCell），返回所有时间步的输出序列及网络最终状态。initial_state需传入与cell结构匹配的状态元组，dtype用于指定输出和状态的数据类型。

tf.keras.layers.RNN #

Keras封装的RNN执行层，仅接受单个RNNCell实例（StackedRNNCells是多cell堆叠的封装实例，符合要求），支持time_major输入格式，调用时可传入initial_state指定初始状态，返回所有时间步的输出及最终状态。层初始化时指定dtype可设置默认数据类型。

内容的提问来源于stack exchange，提问作者user23255480