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.