解决Tensorflow sess.run导致的内存溢出问题

更新时间：2020年02月05日 09:57:07 作者：hustwayne

今天小编就为大家分享一篇解决Tensorflow 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")

下面是解决溢出问题的代码（将部分代码放在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")

总结：

在迭代过程中, 在sess.run的for循环中不要加入tensorflow一些op操作，会增加图节点，否则随着迭代的进行，tf的图会越来越大，最终导致溢出；

建议不要使用tf.gfile.FastGFile(image_path, 'rb').read()读入数据（有可能会造成溢出），用opencv之类读取。

以上这篇解决Tensoflow sess.run导致的内存溢出问题就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持脚本之家。

您可能感兴趣的文章:

python 计算t分布的双侧置信区间
这篇文章主要介绍了python 计算t分布的双侧置信区间，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2021-04-04
python实现批量按比例缩放图片效果
这篇文章主要为大家详细介绍了python实现批量按比例缩放图片效果，具有一定的参考价值，感兴趣的小伙伴们可以参考一下
2018-03-03
用Python实现一个简单的多线程TCP服务器的教程
这篇文章主要介绍了用Python实现一个简单的多线程TCP服务器的教程,示例的运行环境为Windows操作系统,需要的朋友可以参考下
2015-05-05
python实现将文件名批量命名为四位数or五位数
这篇文章主要介绍了python实现将文件名批量命名为四位数or五位数问题,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教
2023-08-08
python list语法学习（带例子）
python list语法学习
2013-11-11
python的time模块和datetime模块实例解析
这篇文章主要介绍了python的time模块和datetime模块实例解析,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下
2019-11-11
聊聊python在linux下与windows下导入模块的区别说明
这篇文章主要介绍了聊聊python在linux下与windows下导入模块的区别说明，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2021-03-03
Python字典的基础操作
这篇文章主要介绍了Python字典的基础操作，Python中的字典数据类型和现实中的字典很像，它是以键值对（键和值的组合）的方式把数据组织到一起，可以通过键找到与之对应的值并进行操作，下面来看文章金额提内容吧,需要的朋友可以参考一下
2021-11-11
解决Python 遍历字典时删除元素报异常的问题
下面小编就为大家带来一篇解决Python 遍历字典时删除元素报异常的问题。小编觉得挺不错的，现在就分享给大家，也给大家做个参考。一起跟随小编过来看看吧
2016-09-09
利用Python绘制好看的玫瑰花图
玫瑰花图是节理统计方式之一，方法简便，形象醒目，比较清楚地反映出主要节理的方向，有助于分析区域构造。本文将用Python绘制好看的玫瑰花图，感兴趣的可以了解一下
2022-05-05