使用Keras预测触发OOM内存错误，训练阶段无异常

核心原因 #

训练阶段与预测阶段的内存管理逻辑存在差异，导致看似矛盾的OOM问题：

• 训练时TensorFlow会自动释放反向传播后的中间张量，且训练过程的内存优化策略（如梯度 checkpoint、自动混合精度）会更激进；
• 预测阶段默认保留更多计算图张量，若存在内存碎片、数据加载不当或模型推理模式未正确切换，就会触发OOM，哪怕batch_size=1。

具体排查与修复步骤 #

1. 清理GPU内存碎片
   训练结束后GPU内存可能存在大量碎片，即使总内存充足，也无法分配连续的内存块。在预测前执行以下代码：

   import tensorflow as tf
   tf.keras.backend.clear_session()
   # 启用GPU内存动态增长，避免一次性占满显存
   gpus = tf.config.list_physical_devices('GPU')
   if gpus:
       try:
           tf.config.experimental.set_memory_growth(gpus[0], True)
       except RuntimeError as e:
           print(e)
   

2. 确保测试数据加载方式正确
   不要一次性将所有测试数据加载为大张量传入model.predict，即使batch_size=1，TensorFlow也会尝试将整个数据张量复制到GPU。改用tf.data.Dataset分批加载：

   test_dataset = tf.data.Dataset.from_tensor_slices(test_data).batch(1)
   predictions = model.predict(test_dataset)
   

   同时检查测试数据的预处理逻辑，确保输入形状、数据类型与训练时完全一致（比如训练时输入是(224,224,3)，测试时不能出现更大的尺寸）。

3. 强制模型进入推理模式
   部分层（如Dropout、BatchNormalization）在训练和推理模式下的行为不同，若未正确切换，可能产生额外的内存开销。加载模型后显式设置：

   model = tf.keras.models.load_model('your_model_path.h5')
   model.trainable = False
   # 或者用tf.function包裹预测逻辑，优化计算图
   @tf.function
   def predict_fn(inputs):
       return model(inputs, training=False)
   

4. 逐样本预测并手动释放内存
   若上述方法无效，尝试循环处理每个样本，每次预测后清理张量：

   predictions = []
   for sample in test_data:
       # 扩展维度以匹配模型输入（假设模型输入是(batch_size, ...)）
       pred = model.predict(np.expand_dims(sample, axis=0), batch_size=1)
       predictions.append(pred[0])
       # 手动删除张量并清理会话
       del pred
       tf.keras.backend.clear_session()
   

5. 检查模型输出层与中间层的内存占用
   虽然输出是137K维度的softmax，但OOM由model/output/Softmax发起，说明输入到softmax的logits张量可能异常。检查模型的全连接层输出维度是否正确，是否在预测时产生了超出预期的大张量（比如训练时用了维度压缩，预测时遗漏）。

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