如何降低TensorFlow模型循环执行预测任务时的内存占用

这个内存溢出的核心原因是TensorFlow默认执行逻辑会累积前向传播的张量引用、重复的计算图trace缓存，你目前用的gc和K.clear_session对Hub加载的SavedModel无效，可按以下步骤修复：

• 调用模型时关闭梯度跟踪，避免TensorFlow留存不必要的运算张量
  所有推理过程都放在tf.no_gradient()上下文管理器内，阻止计算图跟踪推理输出的张量，从根源减少内存累积。
• 固定输入尺寸+给模型添加静态输入签名，避免重复trace
  你每次读入的图片尺寸如果有细微差异，TensorFlow会为每次新输入的shape生成独立的计算图trace，大量trace缓存会快速占满内存。你需要先将所有输入图片统一resize到固定尺寸，再将模型封装为指定了输入签名的tf.function，全程仅生成一次trace。
• 仅提取必要的推理结果，释放TensorFlow张量引用
  模型返回的out是关联了计算图的张量字典，不要留存完整的out对象，仅提取你需要的字段并转为numpy数组，完全释放TensorFlow侧的张量内存。
• 删除无效的清退代码
  K.clear_session()仅对Keras构建的模型生效，对Hub加载的SavedModel没有作用，反复调用反而可能产生额外的内存碎片，可直接删除该行。

修复后的参考代码 #

import tensorflow as tf
import gc
import tensorflow_hub as hub

# 加载模型后封装为固定输入签名的tf.function，避免重复trace
detector = hub.load("https://tfhub.dev/tensorflow/ssd_mobilenet_v2/2")
# 输入签名和你实际的输入对齐即可，示例为单张640*640的3通道uint8图片
detect_fn = tf.function(
    detector,
    input_signature=[tf.TensorSpec(shape=(1, 640, 640, 3), dtype=tf.uint8)]
)

for i in range(num_images):
    img = read_img(img_paths[i])
    # 统一resize到固定尺寸，和上面的input_signature对齐
    img = tf.image.resize(img, (640, 640), preserve_aspect_ratio=False)
    img = tf.cast(img[tf.newaxis, ...], tf.uint8)
    # 关闭梯度跟踪再推理
    with tf.no_gradient():
        out = detect_fn(img)
    # 仅提取需要的字段转为numpy，释放tf张量引用，按需保留你需要的字段即可
    res = {
        "boxes": out["detection_boxes"].numpy(),
        "scores": out["detection_scores"].numpy(),
        "classes": out["detection_classes"].numpy()
    }
    # 此处添加你的业务逻辑处理res
    del img, out, res
    gc.collect()

如果仍有小幅内存上涨，可在代码初始化阶段添加tf.config.optimizer.set_jit(False)关闭XLA编译缓存，避免编译缓存累积占用内存。

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