Tensorflow enables serialization and deserialization of models in many ways. This article shows how to use the Protocol Buffer to write and load a Tensorflow session graph. The important thing is to know the names of the input and output layers of the graph.

Saving the graph to ‘.pb’ file

def save_to_protobuf(sess, path):
    graph = sess.graph.as_graph_def()
    with open(path, 'wb') as file:
        file.write(graph.SerializeToString())

save_to_protobuf(sess, 'test.pb')

Reading the graph from ‘.pb’ file

def load_from_protobuf(path):
    tf.reset_default_graph()
    sess = tf.Session()

    with tf.gfile.GFile(path,'rb') as file:
        graph_def = tf.GraphDef()        
        graph_def.ParseFromString(file.read())
    
    sess.graph.as_default()
    tf.import_graph_def(graph_def)
    return sess

sess = load_from_protobuf('test.pb')

Making inference

input_layer = sess.graph.get_tensor_by_name("import/__input_layer_name__:0")
output_layer = sess.graph.get_tensor_by_name("import/__output_layer_name__:0")

sess.run(output_layer, feed_dict={input_layer: __input_data__})