gen_graph_config_header.py

import argparse
import sys, os
import tensorflow as tf
from get_graph_config_param import parse

def get_io_size(model_dir, input_node_name, output_node_name):
    """ Return tuple consisiting of the input size, and output size"""
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # To suppress info and warning log messages from tensorflow
    clear_devices = True
    checkpoint = tf.train.get_checkpoint_state(model_dir)
    input_checkpoint = checkpoint.model_checkpoint_path
    # We start a session using a temporary fresh Graph
    with tf.Session(graph=tf.Graph()) as sess:
        # We import the meta graph in the current default Graph
        saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=clear_devices)

        # We restore the weights
        saver.restore(sess, input_checkpoint)
        graph = tf.get_default_graph()
        x = graph.get_tensor_by_name('{}:0'.format(input_node_name))
        y = graph.get_tensor_by_name('{}:0'.format(output_node_name))
        return x.shape[1], y.shape[1]

def write_header(output_dir, header_filename, input_size, output_size):
    filepath = os.path.join(output_dir, header_filename)
    with open(filepath, 'w') as f:
        f.write("// Generated by gen_graph_config_header.py, DO NOT EDIT\n")
        f.write("#define GRAPH_INPUT_SIZE {}\n".format(input_size))
        f.write("#define GRAPH_OUTPUT_SIZE {}\n".format(output_size))

def generate_header(model_dir, output_dir, header_filename, config_name):
    input_node_name = parse(model_dir, True, False)
    output_node_name = parse(model_dir, False, True)
    input_size, output_size = get_io_size(model_dir, input_node_name, output_node_name)
    write_header(output_dir, header_filename, input_size, output_size)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("model_dir", help="Directory where the Tensorflow checkpoints are located.")
    parser.add_argument("output_dir", help="Output directory where the file should be saved.")
    parser.add_argument("--header-file-name", help="Name of the header file", default="graph_dim.h")
    parser.add_argument("--config-name", help="Name of the tf2xla proto config file used for tfcompile.", default="tf2xla.config.pbtxt")
    args = parser.parse_args()
    generate_header(args.model_dir, args.output_dir, args.header_file_name, args.config_name)