load.py

"""
Load SavedModel

Output graphdef and checkpoint files
"""

import tensorflow as tf
import argparse
import sys

FLAGS = None

def main(_):
    with tf.Session() as sess:
        print("Loading SavedModel")
        tf.saved_model.loader.load(sess, [FLAGS.tag_name], FLAGS.saved_model_dir)
        # everything is now in sess.graph

        print("Saving graph def to protobuf")
        tf.train.write_graph(sess.graph_def, '.', FLAGS.output_graph)

        saver = tf.train.Saver()
        print("Saving checkpoint")
        saver.save(sess, FLAGS.checkpoint_dir)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--saved_model_dir', type=str, help='Directory of SavedModel, there should be a .pb or .pbtxt file there')
    parser.add_argument('--tag_name', type=str, help='tag name of SavedModel')
    parser.add_argument('--output_graph', type=str, help='Name of output graph, i.e. graph.pbtxt (.pbtxt)')
    parser.add_argument('--checkpoint_dir', type=str, help='Checkpoint directory')

    FLAGS, unparsed = parser.parse_known_args()
    tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)