TensorFlow训练LSTM时遇ResourceExhausted Error或OOM问题求助

作为天天跟TensorFlow内存问题死磕的开发者，我整理了几个亲测有效的解决方案，你可以根据自己的硬件和需求来尝试：

1. 先砍批量大小（最直接的操作） #

OOM大概率是一次性塞给GPU的样本太多了，直接把batch_size往小了调——比如从64降到32，甚至16：

# 如果你用model.fit
model.fit(X_train, y_train, batch_size=16, epochs=10, ...)

# 如果你用tf.data.Dataset
train_dataset = train_dataset.batch(16)

别担心，小批量训练反而有时候能带来更好的泛化性，只是训练速度会慢一点。

2. 给LSTM模型“瘦身” #

LSTM的单元数和层数是内存老虎：

• 减少每层的单元数量，比如把256改成128或者64：

  tf.keras.layers.LSTM(128, return_sequences=True)
  

• 砍掉多余的LSTM层，比如从3层减到2层，先保证模型能跑起来，再慢慢加复杂度调效果。

3. 开启混合精度训练 #

TensorFlow的混合精度能让部分计算用半精度（float16），内存占用直接砍半，而且对精度影响极小：

from tensorflow.keras.mixed_precision import set_global_policy
set_global_policy('mixed_float16')

注意最后输出层最好指定dtype='float32'，避免最后的结果精度丢失：

tf.keras.layers.Dense(num_classes, dtype='float32')

4. 优化数据加载逻辑 #

别一次性把所有数据都塞进内存，用tf.data.Dataset来高效处理：

train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
# 打乱+分批+预取，让数据加载和模型训练并行
train_dataset = train_dataset.shuffle(len(X_train)).batch(16).prefetch(tf.data.AUTOTUNE)

另外检查你的时序序列长度，如果太长的话，试试截断成更短的片段，或者用滑动窗口生成短序列，每个样本占的内存就少了。

5. 让TensorFlow按需占用GPU内存 #

默认TensorFlow会把GPU内存全占了，你可以设置成动态增长，用多少拿多少：

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print(f"可用物理GPU：{len(gpus)}，逻辑GPU：{len(tf.config.list_logical_devices('GPU'))}")
    except RuntimeError as e:
        print(f"设置GPU内存增长失败：{e}")

6. 排查内存泄漏问题 #

有时候OOM是因为内存泄漏：

• 每次训练前清理一下会话，避免残留的张量占内存：

  tf.keras.backend.clear_session()
  

• 别在训练循环里重复创建模型或者张量，这些都会偷偷占内存。

7. 临时切换到CPU训练（应急方案） #

如果GPU实在扛不住，先切到CPU跑通模型逻辑，验证没问题了再想GPU的事：

with tf.device('/CPU:0'):
    model.fit(X_train, y_train, epochs=10, ...)

缺点就是速度慢，但总比卡着不动强。

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