Permalink
Find file
c784bf6 May 11, 2016
30 lines (19 sloc) 660 Bytes
from convert import print_prob, load_image, checkpoint_fn, meta_fn
import tensorflow as tf
layers = 50
img = load_image("data/cat.jpg")
sess = tf.Session()
new_saver = tf.train.import_meta_graph(meta_fn(layers))
new_saver.restore(sess, checkpoint_fn(layers))
graph = tf.get_default_graph()
prob_tensor = graph.get_tensor_by_name("prob:0")
images = graph.get_tensor_by_name("images:0")
for op in graph.get_operations():
print op.name
#init = tf.initialize_all_variables()
#sess.run(init)
print "graph restored"
batch = img.reshape((1, 224, 224, 3))
feed_dict = {images: batch}
prob = sess.run(prob_tensor, feed_dict=feed_dict)
print_prob(prob[0])