如何为TensorFlow自定义LSTM类编写call函数？

一、核心逻辑说明 #

你当前定义的是「CNN提取单帧空间特征 + LSTM提取时序特征」的组合模型，通常用于视频分类、时序图像识别类任务，默认输入为带时序维度的图像序列，形状为(批量大小, 时间步长, 图像高度, 图像宽度, 通道数)。

二、LSTMModel 完整实现（含call方法） #

你已经在__init__中完成了所有层的初始化，只需要在call方法中实现「单帧特征提取→时序特征拼接→LSTM时序推理→分类输出」的数据流逻辑即可，完整代码如下：

import tensorflow as tf

class LSTMModel(tf.keras.Model):
    def __init__(self, CNN_model, num_classes):
        super().__init__()
        # 传入预定义的CNN特征提取器
        self.cnn_model = CNN_model
        # LSTM层：units=64为隐层维度，return_state=True会额外返回隐状态h和细胞状态c
        # 若不需要中间状态可将return_state改为False，代码会更简洁
        self.lstm = tf.keras.layers.LSTM(units=64, return_state=True, dropout=0.3)
        # 分类头
        self.dense = tf.keras.layers.Dense(num_classes, activation="softmax")

    def call(self, inputs, training=False):
        batch_size, timesteps, h, w, c = tf.shape(inputs)
        # 1. 合并批量和时序维度，一次性通过CNN提取所有帧特征
        # 变形后形状：(batch_size * timesteps, h, w, c)，符合CNN输入要求
        cnn_input = tf.reshape(inputs, (-1, h, w, c))
        # 输出形状：(batch_size * timesteps, 1024)，对应CNN最后全连接层的输出维度
        cnn_output = self.cnn_model(cnn_input, training=training)
        # 2. 恢复时序维度，得到时序特征序列，形状：(batch_size, timesteps, 1024)
        lstm_input = tf.reshape(cnn_output, (batch_size, timesteps, -1))
        # 3. 传入LSTM层，return_state=True时返回三个值：最后时刻输出、隐状态h、细胞状态c
        # 若不需要状态可直接写为 lstm_output = self.lstm(lstm_input, training=training)
        lstm_output, state_h, state_c = self.lstm(lstm_input, training=training)
        # 4. 传入全连接层得到分类结果
        output = self.dense(lstm_output)
        return output

参数说明：新增的training参数用于控制dropout、BatchNormalization等训练/推理阶段行为不同的层正常工作，属于TensorFlow自定义模型的标准写法。

三、可选优化：修正现有CNN类的问题 #

你当前写的generic_vns_function存在两处明显问题，可同步调整：

1. 卷积层的第一个参数为滤波器数量，原代码误用了循环索引作为参数，会导致卷积层参数异常
2. 传入的layer_units参数未被使用，可复用为全连接层维度控制参数
   修正后的代码如下：

class generic_vns_function(tf.keras.Model):
    def __init__(self, input_shape, layers_filters, dense_units=1024): 
        super().__init__() 
        self.convolutions = []
        # layers_filters为每一层卷积的滤波器数量列表，例如[32,64,128]
        for filter_num in layers_filters:
            self.convolutions.append(tf.keras.layers.Conv2D(filter_num, 3, padding="same", 
                input_shape=input_shape, activation="relu"))
            # 可选优化：每层卷积后加池化，避免特征尺寸过大，效果优于所有卷积做完再加一层池化
            self.convolutions.append(tf.keras.layers.MaxPooling2D((2,2)))
        
        self.flatten = tf.keras.layers.Flatten()
        self.dense1 = tf.keras.layers.Dense(dense_units, activation="relu")

    def call(self, inputs, training=False):
        x = inputs
        for layer in self.convolutions:
            x = layer(x)
        x = self.flatten(x)
        x = self.dense1(x)
        return x

四、使用示例 #

# 初始化CNN：单帧图像尺寸为64*64、3通道，三层卷积滤波器数量分别为32、64、128
input_shape = (64,64,3)
cnn = generic_vns_function(input_shape, layers_filters=[32,64,128], dense_units=1024)
# 初始化LSTM模型，10分类任务
lstm_model = LSTMModel(cnn, num_classes=10)
# 测试输入：批量大小2、时序长度16、单帧尺寸64*64*3
test_input = tf.random.normal((2, 16, 64, 64, 3))
output = lstm_model(test_input)
print(output.shape) # 输出为(2,10)，符合预期

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