tensorflow frozen graph时对lookup table的处理
https://github.com/tensorflow/tensorflow/issues/8665
两个解决办法:
一个是把 init_all_tables 加到导出图里
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()
另外一种是用tables_initializer配合get_collection, 但需要用meta_graph
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()
Comments