TensorFlow自定义目标检测模型推理内存泄漏无法释放问题

核心根因 #

基于SSD ResNet101 FPN导出的SavedModel检测模型属于混合执行模式，直接调用推理时存在两个内存持有问题：

1. Eager模式下每次推理会动态生成临时计算节点挂载到全局默认计算图，同时为卷积等算子预分配工作缓存，Python原生GC和单独调用tf.keras.backend.clear_session()无法触达这部分被TF runtime持有的内存
2. 推理输出的结果是TF原生张量类型，会被计算图持有引用，不会随Python变量回收自动释放

可落地修复方案 #

• 编译静态推理图，输出结果即时转numpy数组切断引用
  这一步可以解决80%以上的单次推理内存残留问题，代码参考：

  # 模型只加载一次，禁止在推理循环里重复加载
  detect_fn = tf.saved_model.load(visa_icon_model)
  # 将推理函数编译为静态图，固定输入格式，避免每次推理动态生成图节点
  input_signature = [tf.TensorSpec(shape=(1, 模型输入高, 模型输入宽, 3), dtype=tf.uint8)]
  detect_fn = tf.function(detect_fn, input_signature=input_signature)
  # 首次推理前用dummy输入预热，提前分配算子工作内存，避免首次推理的内存波动被误判为泄漏
  _ = detect_fn(tf.zeros((1, 模型输入高, 模型输入宽, 3), dtype=tf.uint8))
  
  # 正式推理
  visa_icon_detections = detect_fn(input_tensor)
  # 关键：立刻将输出张量转为普通numpy数组，解除TF runtime对结果内存的持有
  visa_icon_detections = {key: val.numpy() for key, val in visa_icon_detections.items()}
  

• 补全TF静态图重置逻辑，不要只调用gc.collect()
  之前用的清理方法没覆盖静态图节点的回收，单批次推理完成后按顺序调用以下清理逻辑：

  import gc
  # 先删除所有指向TF张量、模型实例的Python变量
  del visa_icon_detections
  del input_tensor
  # 依次清理Keras会话、重置默认计算图，最后触发Python GC
  tf.keras.backend.clear_session()
  tf.compat.v1.reset_default_graph()
  gc.collect()
  

• 批量推理场景用子进程隔离兜底
  TF Object Detection API导出的模型存在部分历史遗留的图缓存问题，如果同进程内反复推理还是存在内存线性上涨，直接用子进程跑推理逻辑，推理完成后销毁子进程，操作系统会直接回收所有占用内存，不存在残留问题。子进程推理的参考实现：

  from multiprocessing import Process, Queue
  
  def detection_worker(model_path, input_queue, output_queue):
      import tensorflow as tf
      # 子进程内单独加载模型、编译静态图、预热
      local_detect_fn = tf.saved_model.load(model_path)
      input_signature = [tf.TensorSpec(shape=(1, 640, 640, 3), dtype=tf.uint8)]
      local_detect_fn = tf.function(local_detect_fn, input_signature=input_signature)
      _ = local_detect_fn(tf.zeros((1, 640, 640, 3), dtype=tf.uint8))
      while True:
          input_data = input_queue.get()
          if input_data is None: # 收到终止信号退出
              break
          raw_res = local_detect_fn(input_data)
          parsed_res = {k: v.numpy() for k, v in raw_res.items()}
          output_queue.put(parsed_res)
  
  # 主进程调用逻辑
  if __name__ == "__main__":
      in_q, out_q = Queue(), Queue()
      det_process = Process(target=detection_worker, args=(visa_icon_model, in_q, out_q))
      det_process.start()
      
      # 送入待推理数据
      in_q.put(input_tensor)
      final_result = out_q.get()
      
      # 推理全部完成后销毁子进程，内存完全释放
      in_q.put(None)
      det_process.join()
  

• 预处理逻辑尽量移出TF计算图
  不要把图片解码、缩放、归一化等预处理操作写成TF算子放在推理流程里，尽量用OpenCV、numpy在CPU上完成所有预处理，最终输出符合模型输入要求的numpy数组，再转为TF张量喂入模型，减少静态图里的动态节点数量，降低内存残留概率。

注意 #

观测到的单次推理新增191.8MiB内存有一部分是卷积算子预分配的工作空间缓存，属于正常内存占用，只要后续连续推理时内存不会持续线性增长，就不属于内存泄漏。

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