Tensorflow GraphDef pb 文件读和写 （binary format text format, ）


import tensorflow as tf

import os.path

from tensorflow.python.platform import gfile

from google.protobuf import text_format

FLAGS = tf.app.flags.FLAGs

#Input Graph model file location

tf.app.flags.DEFINE_string('model_dir', '/home/eunice/data/test_code', """Paht to classify_image_graph_def.pb""")

#Output Graph model protobuf as text format & binary format

tf.app.flags.DEFINE_string('output_graph_txt', '/home/eunice/data/test_code/output_graph.pbtxt', """pbtxt""")

tf.app.flags.DEFINE_string('output_graph_pb', '/home/eunice/data/test_code/output_graph.pb', """pb""")

def convert_pb_to_pbtxt():

  model_filename = os.path.join(FLAGS.model_dir, 'classify_image_graph_def.pb')

  with gfile.FastGFile(model_filename, 'rb') as f:

   graph_def = tf.GraphDef()        #创建一个GraphDef

   graph_def.ParseFromString(f.read())   #ParseFromString(), reading message from protocol buffer binary fomat

  

  with gfile.FastGFile(FLAGS.output_graph_txt, 'wb') as f:

   f.write(text_format.MessageToString(graph_def))   # MessageToString(message, as_utf8=False, as_one_line=False)  Convert protobuf message to text format

  

#  with gfile.FastGFile(FLAGS.output_graph_pb, 'wb') as f:

#   f.write(graph_def.SerializeToString())  #serializes the message and returns it as a string. Note that the bytes are binary, not text; we only use the str type as a convenient container.