如何在自定义TensorFlow模型中实现可变数量的堆叠LSTM层

该需求完全可以实现，核心修改点是通过tf.keras.layers.LayerList容器存储动态创建的多层LSTM，同时注意非最后一层LSTM需要开启return_sequences=True参数，保证下一层LSTM可以接收序列格式的输入。

修改后的完整代码 #

import tensorflow as tf

class CustomizedLSTM(tf.keras.Model):
    def __init__(self, num_hidden_layers, vocab_size, lstm_units=256):
        super(CustomizedLSTM, self).__init__()
        self.embedding = tf.keras.layers.Embedding(vocab_size, 300)
        # 用LayerList存储堆叠的LSTM层
        self.lstm_layers = tf.keras.layers.LayerList()
        for i in range(num_hidden_layers):
            # 最后一层LSTM不需要返回序列，前面的所有层都需要返回完整序列
            return_seq = False if i == num_hidden_layers - 1 else True
            self.lstm_layers.append(
                tf.keras.layers.LSTM(lstm_units, activation="relu", return_sequences=return_seq)
            )
        self.first_dense = tf.keras.layers.Dense(64, activation="relu")
        self.classification_layer = tf.keras.layers.Dense(1, activation="sigmoid")

    def call(self, inputs):
        x = self.embedding(inputs)
        # 依次调用所有LSTM层
        for lstm_layer in self.lstm_layers:
            x = lstm_layer(x)
        x = self.first_dense(x)
        return self.classification_layer(x)

关键说明 #

• 关于return_sequences参数：LSTM层默认仅返回最后一个时间步的输出（形状为[batch_size, units]的二维张量），而堆叠结构中下层LSTM需要接收三维的序列输入[batch_size, timesteps, units]，因此除最后一层LSTM外，其余层必须开启return_sequences=True。
• 关于LayerList容器：不要使用普通Python列表存储Keras层，普通列表内的层参数不会被TensorFlow自动追踪，会导致训练时参数缺失、模型无法正常保存加载的问题，LayerList是Keras专门用于存储多层结构的容器，会自动注册所有层的参数。
• 扩展性优化：如果需要为不同LSTM层设置不同的单元数、激活函数等参数，可以新增参数如lstm_unit_list传入每层的单元配置，只需保证列表长度和num_hidden_layers一致即可。

使用示例 #

比如要创建3层堆叠LSTM的模型，调用方式如下：

# 3层LSTM，词表大小10000
model = CustomizedLSTM(num_hidden_layers=3, vocab_size=10000)
# 测试输入 shape为[batch_size, seq_len]
test_input = tf.random.uniform((16, 50), maxval=10000, dtype=tf.int32)
output = model(test_input)
print(output.shape) # 输出应为(16, 1)

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