如何使用predict_on_batch避免DataGenerator场景下GPU内存不足

核心问题分析 #

你遇到的两个问题本质关联：

• 显存不足：大数据集+高维度输出导致submodel.predict(train_ds)一次性加载过多数据到GPU
• 输出差异：遍历单样本时未切换模型到推理模式，导致BatchNormalization等依赖训练/推理状态的层行为异常（比如用单样本的均值方差而非训练阶段累积的移动统计量）

解决方案：分批推理+强制推理模式 #

以下步骤可同时解决显存问题，并保证输出与submodel.predict(train_ds)完全一致：

1. 切换子模型到推理模式

预测前必须确保模型进入推理状态，让BatchNormalization、Dropout等层使用训练阶段的累积参数，而非当前batch的临时统计量：

import tensorflow as tf
import numpy as np

# 锁定子模型为推理模式
submodel.trainable = False
# 或用上下文管理器更严谨（适配TensorFlow 2.x）
# with tf.keras.backend.learning_phase(0):

2. 用predict_on_batch分批处理数据集

直接遍历数据集的批次，逐批处理并收集结果，避免一次性加载全部数据：

# 初始化存储latent数据的列表
latent_outputs = []

# 遍历训练数据集的每个batch
for batch_data in train_ds:
    # 逐批预测，自动适配当前batch size
    batch_latent = submodel.predict_on_batch(batch_data)
    latent_outputs.append(batch_latent)

# 合并所有批次结果，顺序与原predict输出完全一致
full_latent_data = np.concatenate(latent_outputs, axis=0)

3. 验证输出一致性

取少量样本对比两种方式的输出，确保仅存在浮点精度级别的差异：

# 取前100个样本对比
small_ds = train_ds.take(100).unbatch().batch(100)
predict_result = submodel.predict(small_ds)
batch_result = np.concatenate([submodel.predict_on_batch(b) for b in small_ds], axis=0)

# 验证一致性（允许1e-6的浮点误差）
assert np.allclose(predict_result, batch_result, atol=1e-6), "输出不一致"

关键注意事项 #

• 若train_ds是自定义DataGenerator，只需保证每次迭代返回标准batch数据，代码逻辑完全通用
• 若显存仍紧张，可进一步缩小数据集的batch size（比如从64改为32），不影响最终输出一致性
• 不要手动调用模型的__call__方法（如submodel(batch_data)），这会默认使用训练模式，需显式设置training=False：submodel(batch_data, training=False)

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