解决TensorFlow sess.run内存溢出：避免循环内动态构建计算图

问题现象：sess.run导致的内存溢出与性能下降

在调用模型进行批量测试时，代码出现了内存溢出错误。起初怀疑是数据读取方式导致的问题，尝试了tf、PIL和cv等多种图片读取方式，结果发现读取速度越来越慢，内存占用持续飙升。经过调试定位，问题核心在于sess.run的执行过程：随着for循环次数的增加，程序运行速度显著下降，内存消耗急剧上升。

  # Creates graph from saved GraphDef
  create_graph(pb_path)
 
  # Init tf Session
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True
  sess = tf.Session(config=config)
  init = tf.global_variables_initializer()
  sess.run(init)
 
 
  input_image_tensor = sess.graph.get_tensor_by_name("create_inputs/batch:0") 
  output_tensor_name = sess.graph.get_tensor_by_name("conv6/out_1:0") 
 
 
  for filename in os.listdir(image_dir):
    image_path = os.path.join(image_dir, filename)
 
    start = time.time()
    image_data = cv2.imread(image_path)
    image_data = cv2.resize(image_data, (w, h))
    image_data_1 = image_data - IMG_MEAN
    input_image = np.expand_dims(image_data_1, 0)
 
    raw_output_up = tf.image.resize_bilinear(output_tensor_name, size=[h, w], align_corners=True) 
    raw_output_up = tf.argmax(raw_output_up, axis=3)
    
 
    predict_img = sess.run(raw_output_up, feed_dict={input_image_tensor: input_image})    # 1，height，width
    predict_img = np.squeeze(predict_img)   # height， width 
 
    voc_palette = visual.make_palette(3)
    masked_im = visual.vis_seg(image_data, predict_img, voc_palette)
    cv2.imwrite("%s_pred.png" % (save_dir + filename.split(".")[0]), masked_im)
 
 
    print(time.time() - start)
 
  print(">>>>>>Done")

以下是解决该溢出问题的代码示例，主要改进点是将部分TensorFlow操作代码移至for循环外部执行。

  # Creates graph from saved GraphDef
  create_graph(pb_path)
 
  # Init tf Session
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True
  sess = tf.Session(config=config)
  init = tf.global_variables_initializer()
  sess.run(init)
 
  input_image_tensor = sess.graph.get_tensor_by_name("create_inputs/batch:0") 
  output_tensor_name = sess.graph.get_tensor_by_name("conv6/out_1:0") 
  
##############################################################################################################
  raw_output_up = tf.image.resize_bilinear(output_tensor_name, size=[h, w], align_corners=True) 
  raw_output_up = tf.argmax(raw_output_up, axis=3)
##############################################################################################################
 
  for filename in os.listdir(image_dir):
    image_path = os.path.join(image_dir, filename)
 
    start = time.time()
    image_data = cv2.imread(image_path)
    image_data = cv2.resize(image_data, (w, h))
    image_data_1 = image_data - IMG_MEAN
    input_image = np.expand_dims(image_data_1, 0)
    
    predict_img = sess.run(raw_output_up, feed_dict={input_image_tensor: input_image})    # 1，height，width
    predict_img = np.squeeze(predict_img)   # height， width 
 
    voc_palette = visual.make_palette(3)
    masked_im = visual.vis_seg(image_data, predict_img, voc_palette)
    cv2.imwrite("%s_pred.png" % (save_dir + filename.split(".")[0]), masked_im)
    print(time.time() - start)
 
  print(">>>>>>Done")

总结与优化建议：

核心原因：避免循环内重复添加Op节点

在迭代过程中，务必避免在sess.run的for循环内部添加新的TensorFlow op操作。如果在循环中动态添加op，会导致TensorFlow的计算图节点不断累积。随着迭代进行，计算图体积越来越大，最终引发内存溢出错误。

数据读取建议：替换tf.gfile.FastGFile

此外，建议避免使用tf.gfile.FastGFile(image_path, 'rb').read()的方式读取数据，这种方式可能导致内存溢出风险。推荐使用OpenCV等库进行图像读取，以提高稳定性和效率。

以上关于解决TensorFlow sess.run导致内存溢出问题的内容分享完毕，希望能对读者有所帮助。