克服Graphdef在张量流中不能大于2GB

Question

我使用tensorflow的imageNet训练模型来提取最后一个池层的特征作为新图像数据集的表示向量.

该模型预测新图像如下:

python classify_image.py --image_file new_image.jpeg 

我编辑了主要功能,以便我可以获取图像文件夹并立即返回所有图像的预测,并将特征向量写入csv文件.我是这样做的:

def main(_):
  maybe_download_and_extract()
  #image = (FLAGS.image_file if FLAGS.image_file else
  #         os.path.join(FLAGS.model_dir, 'cropped_panda.jpg'))
  #edit to take a directory of image files instead of a one file
  if FLAGS.data_folder:
    images_folder=FLAGS.data_folder
    list_of_images = os.listdir(images_folder)
  else: 
    raise ValueError("Please specify image folder")

  with open("feature_data.csv", "wb") as f:
    feature_writer = csv.writer(f, delimiter='|')

    for image in list_of_images:
      print(image) 
      current_features = run_inference_on_image(images_folder+"/"+image)
      feature_writer.writerow([image]+current_features)

它适用于大约21张图像但随后因以下错误而崩溃:

  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/framework/ops.py", line 1912, in as_graph_def
    raise ValueError("GraphDef cannot be larger than 2GB.")
ValueError: GraphDef cannot be larger than 2GB.

我想通过调用该方法run_inference_on_image(images_folder+"/"+image)将覆盖先前的图像数据以仅考虑新的图像数据,这似乎不是这种情况.如何解决这个问题？

Answer 1

这里的问题是每次调用都会run_inference_on_image() 将节点添加到同一个图形中,最终会超出最大大小.至少有两种方法可以解决这个问题:

1. 在容易,但缓慢的方法是使用一个不同的默认图形每次调用run_inference_on_image():

   for image in list_of_images:
     # ...
     with tf.Graph().as_default():
       current_features = run_inference_on_image(images_folder+"/"+image)
     # ...
   

2. 将更多的参与,但更有效的方法是修改run_inference_on_image()对多个图像运行.重新定位for循环以围绕此sess.run()调用,您将不再需要在每次调用时重建整个模型,这样可以更快地处理每个图像.