Keras LSTM手动遍历可变序列长度的解决方案咨询

核心问题在于：当输入是可变长度序列时，TensorFlow静态图中时间轴的维度是None，Python原生的range()无法处理这个动态值，必须改用TensorFlow提供的动态循环API来实现遍历逻辑。同时原代码中初始状态固定为batch_size=1的写法也不够灵活，需要适配动态batch大小。

下面提供两种可行的解决方案：

方法一：使用tf.while_loop实现动态循环 #

tf.while_loop是TensorFlow原生的动态循环API，能处理动态变化的序列长度：

import tensorflow as tf
import numpy as np

class MyExample(tf.keras.Model):

    def __init__(self, int_dim, **kwargs):
        super(MyExample, self).__init__(**kwargs)
        self.int_dim = int_dim
        self.lstm = tf.keras.layers.LSTMCell(self.int_dim)
        self.d2 = tf.keras.layers.Dense(self.int_dim)

    def call(self, inputs):
        # 获取动态的batch size和序列长度
        batch_size = tf.shape(inputs)[0]
        seq_len = tf.shape(inputs)[1]
        
        # 动态生成初始状态（适配任意batch size）
        initial_state = (
            tf.zeros((batch_size, self.int_dim)),
            tf.zeros((batch_size, self.int_dim))
        )
        
        # 定义循环的初始变量：当前时间步、当前状态、输出列表
        def loop_body(t, states, outputs):
            # 提取当前时间步的输入
            current_input = inputs[:, t, :]
            lstm_out, new_states = self.lstm(current_input, states)
            d2_out = self.d2(lstm_out)
            # 更新循环变量：时间步+1，新状态，输出列表追加当前结果
            return t + 1, new_states, outputs.write(t, d2_out)
        
        # 初始化输出张量数组（用于存储每一步的结果）
        output_array = tf.TensorArray(
            dtype=tf.float32, size=seq_len, dynamic_size=False
        )
        
        # 执行循环
        _, final_states, final_outputs = tf.while_loop(
            cond=lambda t, *_: t < seq_len,
            body=loop_body,
            loop_vars=(0, initial_state, output_array)
        )
        
        # 将TensorArray转换为常规张量，并恢复时间轴维度
        output_stack = final_outputs.stack()
        # 转置回到 [batch, seq_len, dim] 的形状
        output_stack = tf.transpose(output_stack, [1, 0, 2])
        return output_stack

def generator():
    while True:
      seq_len = np.random.randint(2, 10)
      X = tf.random.uniform((1, seq_len, 5))
      Y = tf.random.uniform((1, seq_len, 5))
      yield X, Y

model = MyExample(5)
model.compile('adam', 'BinaryCrossentropy')
# 测试运行，限制步数避免无限循环
model.fit(generator(), batch_size=1, steps_per_epoch=10)

关键改动说明： #

• 用tf.shape(inputs)[0]和tf.shape(inputs)[1]获取动态维度值，替代静态的inputs.shape（静态维度在可变长度场景下为None）
• 使用tf.TensorArray存储循环中的输出，避免Python列表的静态限制
• 通过tf.while_loop实现动态遍历，循环条件基于动态的序列长度
• 初始状态根据输入的batch size动态生成，适配任意批量大小

方法二：使用tf.scan简化带状态的遍历逻辑 #

如果需要在遍历过程中传递状态（比如LSTM的隐藏状态），tf.scan是更简洁的选择，它会自动遍历指定维度并保留状态传递：

import tensorflow as tf
import numpy as np

class MyExample(tf.keras.Model):

    def __init__(self, int_dim, **kwargs):
        super(MyExample, self).__init__(**kwargs)
        self.int_dim = int_dim
        self.lstm = tf.keras.layers.LSTMCell(self.int_dim)
        self.d2 = tf.keras.layers.Dense(self.int_dim)

    def call(self, inputs):
        batch_size = tf.shape(inputs)[0]
        initial_state = (
            tf.zeros((batch_size, self.int_dim)),
            tf.zeros((batch_size, self.int_dim))
        )
        
        # 定义单步处理函数，tf.scan会遍历时间轴
        def process_step(states, input_step):
            lstm_out, new_states = self.lstm(input_step, states)
            d2_out = self.d2(lstm_out)
            return new_states, d2_out
        
        # 转置输入为 [seq_len, batch, dim] 适配scan的遍历轴
        transposed_inputs = tf.transpose(inputs, [1, 0, 2])
        
        # 执行带状态的遍历
        _, outputs = tf.scan(
            fn=process_step,
            elems=transposed_inputs,
            initializer=initial_state
        )
        
        # 转置回原形状 [batch, seq_len, dim]
        output_stack = tf.transpose(outputs, [1, 0, 2])
        return output_stack

# 生成器和训练代码同前
def generator():
    while True:
      seq_len = np.random.randint(2, 10)
      X = tf.random.uniform((1, seq_len, 5))
      Y = tf.random.uniform((1, seq_len, 5))
      yield X, Y

model = MyExample(5)
model.compile('adam', 'BinaryCrossentropy')
model.fit(generator(), batch_size=1, steps_per_epoch=10)

说明： #

tf.scan专门用于处理需要状态传递的序列遍历场景，相比tf.while_loop代码更简洁，无需手动管理循环变量和输出存储。

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