tensorflow: freeze_graph not initializing tables

Adding init_all_tables to the list of names to export fixes this issue.

import os

import tensorflow as tf
from tensorflow.python.framework.graph_util import convert_variables_to_constants
from tensorflow.python.ops.lookup_ops import HashTable, KeyValueTensorInitializer

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

OUTPUT_FOLDER = '/tmp'
OUTPUT_NAME = 'hash_table.pb'
OUTPUT_NAMES = ['graph/output', 'init_all_tables']


def build_graph():
    d = {'a': 1, 'b': 2, 'c': 3, 'd': 4}
    init = KeyValueTensorInitializer(d.keys(), d.values())
    hash_table = HashTable(init, default_value=-1)
    data = tf.placeholder(tf.string, (None,), name='data')
    values = hash_table.lookup(data)
    output = tf.identity(values * 2, 'output')


def freeze_graph():
    with tf.Graph().as_default() as graph:
        with tf.name_scope('graph'):
            build_graph()

        with tf.Session(graph=graph) as sess:
            sess.run(tf.tables_initializer())
            print sess.run('graph/output:0', feed_dict={'graph/data:0': ['a', 'b', 'c', 'd', 'e']})
            frozen_graph = convert_variables_to_constants(sess, sess.graph_def, OUTPUT_NAMES)
            tf.train.write_graph(frozen_graph, OUTPUT_FOLDER, OUTPUT_NAME, as_text=False)


def load_frozen_graph():
    with open(os.path.join(OUTPUT_FOLDER, OUTPUT_NAME), 'rb') as f:
        output_graph_def = tf.GraphDef()
        output_graph_def.ParseFromString(f.read())

    with tf.Graph().as_default() as graph:
        tf.import_graph_def(output_graph_def, name='')
        with tf.Session(graph=graph) as sess:
            try:
                sess.run(graph.get_operation_by_name('init_all_tables'))
            except KeyError:
                pass
            print sess.run('graph/output:0', feed_dict={'graph/data:0': ['a', 'b', 'c', 'd', 'e']})


if __name__ == '__main__':
    freeze_graph()
    load_frozen_graph()

The call to extract_sub_graph inside convert_variables_to_constants prunes out this op and its descendants (keys, values) if you don’t include init_all_tables in output_node_names. I don’t like the idea of running an initializer op during inference and having a try/except seems hacky to me. Is there another way to do this?

@jkiske 's workaround worked for me!

We also built on top of it to make it work with tf.tables_initializer(), but it requires two other changes:

• OUTPUT_NAMES needs to include the table initialization ops, which can be obtained with tf.get_collection(tf.GraphKeys.TABLE_INITIALIZERS).
• The MetaGraph, instead of the Graph, needs to be what’s exported/imported. This is because tf.tables_initializer() references the tf.GraphKeys.TABLE_INITIALIZERS collection. The Graph does not contain a collection_list, but the MetaGraph does.

So here’s a solution that works for us:

import os

import tensorflow as tf
from tensorflow.python.framework.graph_util import convert_variables_to_constants
from tensorflow.python.ops.lookup_ops import HashTable, KeyValueTensorInitializer

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

OUTPUT_FOLDER = '/tmp'
OUTPUT_NAME = 'hash_table.pb'
OUTPUT_NAMES = ['graph/output']


def build_graph():
    d = {'a': 1, 'b': 2, 'c': 3, 'd': 4}
    init = KeyValueTensorInitializer(d.keys(), d.values())
    hash_table = HashTable(init, default_value=-1)
    data = tf.placeholder(tf.string, (None,), name='data')
    values = hash_table.lookup(data)
    output = tf.identity(values * 2, 'output')


def freeze_graph():
    with tf.Graph().as_default() as graph:
        with tf.name_scope('graph'):
            build_graph()

        with tf.Session(graph=graph) as sess:
            sess.run(tf.tables_initializer())
            print sess.run('graph/output:0', feed_dict={'graph/data:0': ['a', 'b', 'c', 'd', 'e']})
            for table_init_op in tf.get_collection(tf.GraphKeys.TABLE_INITIALIZERS):
                OUTPUT_NAMES.append(table_init_op.name)
            frozen_graph = convert_variables_to_constants(sess, sess.graph_def, OUTPUT_NAMES)
            tf.train.export_meta_graph(
                filename=os.path.join(OUTPUT_FOLDER, OUTPUT_NAME),
                graph_def=frozen_graph,
                collection_list=[tf.GraphKeys.TABLE_INITIALIZERS])


def load_frozen_graph():
    with tf.Graph().as_default() as graph:
        tf.train.import_meta_graph(os.path.join(OUTPUT_FOLDER, OUTPUT_NAME))
        with tf.Session(graph=graph) as sess:
            sess.run(tf.tables_initializer())
            print sess.run('graph/output:0', feed_dict={'graph/data:0': ['a', 'b', 'c', 'd', 'e']})


if __name__ == '__main__':
    freeze_graph()
    load_frozen_graph()

Nagging Assignees @petewarden, @suharshs: It has been 14 days with no activity and this issue has an assignee. Please update the label and/or status accordingly.