TensorFlow图像分类模型预测80万张图片触发OOM内存泄漏问题求助

你遇到的OOM问题不是内存泄漏，是代码逻辑和配置不合理导致的内存持续占用上涨，核心问题如下：

• tf.keras.backend.clear_session() 放置位置错误：你在构造数据集的阶段调用清理逻辑，完全不影响预测阶段的内存回收。TensorFlow预测过程中会持续累积计算图节点、中间张量，无主动清理机制，内存占用会随预测步数线性上升。
• 全量预测结果一次性加载：tf.nn.softmax(model.predict(...)) 会先把80万张图的所有预测输出全部加载到内存再做softmax计算，仅输出张量本身就会占用数GB到数十GB内存，叠加运行时开销直接触发OOM。
• 数据管道配置冗余：
  • 归一化操作放在batch之后，单批次张量体积更大，预加载时占用更多内存
  • 无GPU环境下prefetch(buffer_size=AUTOTUNE)会预加载过多批次到内存，进一步挤占可用空间
• Eager模式未关闭：CPU场景下Eager执行模式不会主动回收每步的中间张量，内存占用上涨速度远高于图模式。

1. 基础配置调整 #

预测启动前先添加全局配置，限制内存占用：

import tensorflow as tf
# 关闭eager执行，用图模式跑预测，大幅降低内存占用
tf.compat.v1.disable_eager_execution()
# 预测前清理残留会话
tf.keras.backend.clear_session()
# 开启XLA编译优化，减少内存占用同时提升预测速度
tf.config.optimizer.set_jit(True)

BATCH_SIZE = 128
# 无GPU场景下prefetch固定为1即可，不需要AUTOTUNE预加载多批
PREFETCH_SIZE = 1

2. 数据管道调整 #

调整map顺序，降低batch级张量体积：

def configure_for_performance(ds):
    ds = ds.batch(BATCH_SIZE)
    ds = ds.prefetch(buffer_size=PREFETCH_SIZE)
    return ds

def decode_img(img):
    img = tf.io.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, [IMG_HEIGHT, IMG_WIDTH])
    # 归一化移到单张图处理阶段，减少batch计算时的内存开销
    img = normalization_sequential_layer(img)
    return img

def process_path(file_path):
    img = tf.io.read_file(file_path)
    img = decode_img(img)
    return img

list_ds = tf.data.Dataset.list_files([filepath1,filepath2,...,filepathN],
                                     shuffle=False)

# 这部分存文件名的逻辑如果不是必须可以删除，减少一次全量迭代
files_list = list()
for files in list_ds.as_numpy_iterator():
    files_list.append(files.decode("utf-8"))

test_ds = list_ds.map(process_path, num_parallel_calls=tf.data.AUTOTUNE)
test_size = test_ds.cardinality().numpy()
test_ds = configure_for_performance(test_ds)

3. 预测逻辑调整 #

分批次处理结果，不要一次性留存全量预测值：

# 如果需要保存预测结果，直接边预测边写磁盘，不要全量存在内存里
with open("predict_result.csv", "w") as f:
    for batch in test_ds:
        batch_score = tf.nn.softmax(model.predict_on_batch(batch)).numpy()
        # 写入当前批次的预测结果，比如存类别、置信度等
        for score in batch_score:
            f.write(f"{score.argmax()},{score.max()}\n")
        # 每批次预测完手动清理张量，释放内存
        del batch_score

如果必须留存全量预测结果，也可以用numpy的memmap把数组存在磁盘上，避免占用内存：

import numpy as np
# 假设分类数为NUM_CLASSES，创建memmap数组存在磁盘，不占内存
score_arr = np.memmap("predict_score.npy", dtype="float32", mode="w+", shape=(test_size, NUM_CLASSES))
idx = 0
for batch in test_ds:
    batch_score = tf.nn.softmax(model.predict_on_batch(batch)).numpy()
    batch_len = len(batch_score)
    score_arr[idx:idx+batch_len] = batch_score
    idx += batch_len
    del batch_score
# 写完后刷到磁盘
score_arr.flush()

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