向ConvLSTM2D模型喂入TFRecord文件时RAM占用过高求助

用TensorFlow搭建了ConvLSTM2D模型，喂入*.tfrecords文件训练时系统内存直接被占满。

模型结构如下：

model = Sequential([
    ConvLSTM2D(64, (3, 3), activation='relu', input_shape=(20, 224, 224, 3), return_sequences=True),
    BatchNormalization(),
    ConvLSTM2D(64, (3, 3), activation='relu', return_sequences=True),
    BatchNormalization(),
    Flatten(),
    Dense(1, activation='sigmoid')
])

训练代码：

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
model.fit(train_input_fn(), steps_per_epoch=5,validation_data=val_input_fn(),epochs=10)

测试环境：2020款M1 MacBook（Jupyter Notebook、Pycharm）、Google Colab

项目背景 #

基于上海数据集（包含打斗和非打斗两类视频），目标是用ConvLSTM2D实现视频分类。训练集有800个打斗视频+800个非打斗视频，我们按250ms间隔提取帧并转成numpy数组，再把所有数组存入TFRecord文件，通过train_input_fn()读取后喂给模型。

数据集结构：

Dataset
    - train
        - Fight      # 800个*.avi文件
        - NonFight   # 800个*.avi文件
    - val
        - Fight      # 200个*.avi文件
        - NonFight   # 200个*.avi文件

已尝试方案 #

• 将batch_size从64降到16
• 把训练集视频数量从800缩减到200
• 减小ConvLSTM2D的滤波器尺寸
• 换用*.mp4格式文件测试
• 移除一层ConvLSTM2D和BatchNormalization层

解决方案 #

1. 修复TFRecord加载逻辑（核心问题） #

内存爆掉大概率是train_input_fn()没有做流式加载，而是一次性把所有数据读进内存了。改成用tf.data.TFRecordDataset做流式处理，配合分批、预取，能直接解决内存占用问题：

def parse_tfrecord_fn(example):
    # 按你存TFRecord时的特征定义写，这里是示例
    feature_description = {
        'frames': tf.io.FixedLenFeature([20, 224, 224, 3], tf.float32),
        'label': tf.io.FixedLenFeature([1], tf.int64)
    }
    example = tf.io.parse_single_example(example, feature_description)
    frames = example['frames']
    # 如果存的是uint8格式的帧，记得转成float32并归一化：frames = tf.cast(frames, tf.float32) / 255.0
    label = tf.cast(example['label'], tf.float32)
    return frames, label

def train_input_fn(batch_size=16):
    # 直接读TFRecord文件，不用一次性加载
    dataset = tf.data.TFRecordDataset(['train.tfrecords'])
    # 并行解析，减少IO等待时间
    dataset = dataset.map(parse_tfrecord_fn, num_parallel_calls=tf.data.AUTOTUNE)
    # 打乱数据，buffer_size不用设成全集大小，100左右足够，省内存
    dataset = dataset.shuffle(buffer_size=100)
    # 分批，drop_remainder避免最后一批数据量不足导致的内存波动
    dataset = dataset.batch(batch_size, drop_remainder=True)
    # 预取下一批数据，让CPU在GPU训练时提前准备，提升效率还不占额外内存
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

另外，存TFRecord时别存已经归一化的float32数组，存uint8格式的原始帧，解析时再做归一化，能大幅减小TFRecord体积和内存占用。

2. 模型端优化 #

• 缩短序列长度：输入序列长度是20，可以降到10或15，ConvLSTM的计算量随序列长度线性增长，能明显降低内存压力。
• 减少滤波器数量：把64改成32，两层ConvLSTM的参数量直接减到原来的1/4，内存占用骤降，后续精度不够再调回来就行。
• 替换Flatten层：用GlobalAveragePooling3D()代替Flatten，Flatten会把(20,224,224,64)转成6400多万个参数的向量，而GlobalAveragePooling3D只输出(64,)的向量，内存差好几个数量级：

model = Sequential([
    ConvLSTM2D(32, (3, 3), activation='relu', input_shape=(15, 224, 224, 3), return_sequences=True),
    BatchNormalization(),
    ConvLSTM2D(32, (3, 3), activation='relu', return_sequences=True),
    BatchNormalization(),
    GlobalAveragePooling3D(),
    Dense(1, activation='sigmoid')
])

3. 环境适配 #

• M1 MacBook：TensorFlow在M1上的内存管理不如x86稳定，建议关闭所有后台无关程序，或者加一行代码开启GPU内存动态增长：

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    tf.config.experimental.set_memory_growth(gpus[0], True)

避免GPU一次性占满内存。

• Google Colab：确保选的是GPU实例，训练前重启内核释放残留内存，同时检查train_input_fn()有没有重复创建数据集导致的内存泄漏。

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