解決Tensorflow sess.run導(dǎo)致的內(nèi)存溢出問題

更新時(shí)間：2020年02月05日 09:57:07 作者：hustwayne

今天小編就為大家分享一篇解決Tensorflow sess.run導(dǎo)致的內(nèi)存溢出問題，具有很好的參考價(jià)值，希望對大家有所幫助。一起跟隨小編過來看看吧

下面是調(diào)用模型進(jìn)行批量測試的代碼(出現(xiàn)溢出)，開始以為導(dǎo)致溢出的原因是數(shù)據(jù)讀入方式問題引起的，用了tf , PIL和cv等方式讀入圖片數(shù)據(jù)，發(fā)現(xiàn)越來越慢，內(nèi)存占用飆升，調(diào)試時(shí)發(fā)現(xiàn)是sess.run這里出了問題（隨著for循環(huán)進(jìn)行速度越來越慢）。

  # 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循環(huán)外）

  # 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")

總結(jié)：

在迭代過程中, 在sess.run的for循環(huán)中不要加入tensorflow一些op操作，會增加圖節(jié)點(diǎn)，否則隨著迭代的進(jìn)行，tf的圖會越來越大，最終導(dǎo)致溢出；

建議不要使用tf.gfile.FastGFile(image_path, 'rb').read()讀入數(shù)據(jù)（有可能會造成溢出），用opencv之類讀取。

以上這篇解決Tensoflow sess.run導(dǎo)致的內(nèi)存溢出問題就是小編分享給大家的全部內(nèi)容了，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。

您可能感興趣的文章: