TensorFlow CPU环境下重复调用DNNClassifier.predict出现OOM错误

嘿，我完全懂你的困惑——明明是在CPU上跑预测，不是训练阶段，按道理预测应该是轻量操作，怎么几千次调用后就爆内存了呢？我之前帮好几个开发者排查过类似的问题，咱们来一步步捋清楚根源和解决办法。

为啥会出现这个问题？ #

你可能以为每次estimator.predict()只是跑个前向传播，但实际上默认情况下，每次调用predict都会创建新的计算图节点、临时张量甚至会话资源，而这些资源不会被及时回收。CPU的内存虽然比GPU显存大，但架不住几千次调用下来，累积的未释放资源把内存占满，最后就触发了ResourceExhaustedError。你提到的错误路径C:\Users\Zvi\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow\python\framework\ops.py，本质就是TensorFlow在尝试分配新内存时，系统已经没空间了。

具体解决办法 #

这里给你几个实用的方案，按优先级排序：

1. 重用预测迭代器，不要每次调用都重新创建

这是最有效的办法，很多人会在循环里每次都调用estimator.predict(input_fn=...)，但其实可以先创建一次迭代器，然后循环获取结果：

# 先定义支持生成样本的输入函数
def predict_input_fn():
    # 返回一个能生成所有待预测样本的数据集
    dataset = tf.data.Dataset.from_generator(
        your_sample_generator,
        output_types=(tf.float32),
        output_shapes=(your_input_shape,)
    )
    return dataset

# 只创建一次预测迭代器
predict_iter = estimator.predict(input_fn=predict_input_fn)

# 循环获取每个样本的预测结果
for _ in range(total_samples):
    prediction = next(predict_iter)
    # 处理你的预测结果

这样就避免了每次调用predict都重新构建计算图和会话，从根源上减少内存累积。

2. 显式管理计算图和会话

如果你的场景需要多次独立调用predict，可以手动指定默认计算图并重用会话，强制TensorFlow不创建新资源：

import tensorflow as tf
import gc

# 提前创建并固定计算图
with tf.Graph().as_default() as fixed_graph:
    # 加载你已经训练好的DNNClassifier
    estimator = tf.estimator.DNNClassifier(
        model_dir="your_trained_model_path",
        feature_columns=your_feature_columns,
        hidden_units=[...]
    )

# 在循环里重用同一个图和会话
with tf.Session(graph=fixed_graph) as sess:
    sess.run(tf.global_variables_initializer())
    for sample in your_thousands_of_samples:
        # 构造单样本输入函数
        def single_sample_input():
            return tf.convert_to_tensor([sample], dtype=tf.float32)
        
        # 传入已有的会话，避免创建新会话
        prediction = next(estimator.predict(input_fn=single_sample_input, session=sess))
        # 处理结果
        # 可选：触发垃圾回收，辅助释放临时资源
        gc.collect()

3. 检查输入函数的内存泄漏

有时候问题不在predict本身，而是你的输入函数每次都会创建大的numpy数组、打开文件但不关闭，或者生成了无法被回收的资源。比如：

• 如果输入函数里每次都读取本地文件，记得用with语句管理文件句柄
• 如果每次都创建大的张量，处理完样本后可以手动赋值为None来释放内存：sample = None

总结 #

核心就是不要让每次predict调用都产生新的计算资源，CPU内存虽然充裕，但架不住几千次的累积。优先试试重用迭代器的方案，这是最直接高效的解决办法。

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