如何用纯TensorFlow 2实现等效Keras的带return_sequence和return_state的LSTM

核心逻辑 #

Keras的LSTM层这两个参数本质是控制输出内容：

• return_sequences=True：返回每个时间步的隐藏状态序列
• return_state=True：返回最后一个时间步的隐藏状态与细胞状态

在纯TF2中，我们可以通过LSTMCell手动构建循环，或者借助tf.nn.dynamic_rnn简化实现，下面是具体方案：

方法一：手动构建循环实现 #

import tensorflow as tf

# 参数定义
n = 64  # 隐藏层维度
ntimesteps = 10
nfeatures = 8
return_sequences = True
return_state = True

# 初始化LSTM细胞
lstm_cell = tf.keras.layers.LSTMCell(n)

# 模拟输入数据（batch_size=32）
inputs = tf.random.normal((32, ntimesteps, nfeatures))
batch_size = tf.shape(inputs)[0]

# 获取初始状态
initial_state = lstm_cell.get_initial_state(inputs=None, batch_size=batch_size, dtype=tf.float32)

# 遍历时间步收集输出
outputs = []
current_state = initial_state
for t in range(ntimesteps):
    step_input = inputs[:, t, :]
    step_output, current_state = lstm_cell(step_input, current_state)
    outputs.append(step_output)

# 整理返回结果
final_output = tf.stack(outputs, axis=1) if return_sequences else step_output
final_states = current_state if return_state else None

# 匹配Keras的输出格式
if return_sequences and return_state:
    result = (final_output, *final_states)
elif return_sequences:
    result = final_output
elif return_state:
    result = (*final_states,)
else:
    result = step_output

方法二：用tf.nn.dynamic_rnn简化实现 #

import tensorflow as tf

# 参数定义
n = 64
ntimesteps = 10
nfeatures = 8
return_sequences = True
return_state = True

lstm_cell = tf.keras.layers.LSTMCell(n)
inputs = tf.random.normal((32, ntimesteps, nfeatures))

# dynamic_rnn自动处理时间步循环
outputs, final_state = tf.nn.dynamic_rnn(
    cell=lstm_cell,
    inputs=inputs,
    dtype=tf.float32
)

# 适配返回格式
result = outputs if return_sequences else outputs[:, -1, :]
if return_state:
    if return_sequences:
        result = (result, *final_state)
    else:
        result = (result, *final_state)

关键说明 #

• return_sequences=True对应保留所有时间步的输出（形状为(batch_size, timesteps, units)），否则仅取最后一步输出（outputs[:, -1, :]，形状为(batch_size, units)）
• return_state=True时，需额外返回LSTM的最终隐藏状态和细胞状态，二者包含在final_state元组中（顺序为隐藏状态、细胞状态）

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