How to load a graph checkpoints (.ckpt) and use SavedModelBuilder to save it as protobuf without declaring any tf.Variables?

mofury picture mofury · Mar 31, 2018 · Viewed 8.6k times · Source

Currently I have resnet_v2_50.ckpt from tensorflow's open sourced pretrained model. I am trying to serve this model in Go cause my backend for my web app is going to be in Go. If I were to create my own model and train it, and then save it. I have no trouble with serving it in Go but I am trying to use pre-trained model to save time on my end.

Here's a simple example of how I save my model

mnist = input_data.read_data_sets(DATA_DIR, one_hot=True)

# Recall that each image is 28x28
x = tf.placeholder(tf.float32, [None, 784], name='imageinput')
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
y = tf.add(tf.matmul(x, W), b)
labels = tf.placeholder(tf.float32, [None, 10])
cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=y))
train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy_loss)

with tf.Session() as sess:
    with tf.device("/cpu:0"):
        sess.run(tf.global_variables_initializer())
        for i in range(1000):
            batch_x, batch_label = mnist.train.next_batch(100)
            loss, _ = sess.run([cross_entropy_loss, train_step], feed_dict={x: batch_x, labels: batch_label})
            print '%d: %f' % (i + 1, loss)

        infer = tf.argmax(y, axis=1, name='infer')
        truth = tf.argmax(labels, axis=1)
        correct_prediction = tf.equal(infer, truth)
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
        print sess.run(accuracy, feed_dict={x: mnist.test.images, labels: mnist.test.labels})

        print 'Time to save the graph!'
        builder = tf.saved_model.builder.SavedModelBuilder('mnist_model')
        builder.add_meta_graph_and_variables(sess, ['serve'])
        builder.save()

I can load it using tensorflow in Go

model, err := tf.LoadSavedModel("./tf_mnist_py/mnist_model", []string{"serve"}, nil)
if err != nil {
    fmt.Printf("Error loading saved model: %s\n", err.Error())
    return
}

defer model.Session.Close()

But now when it comes to pretrained model, I am dealing with ckpt files. One solution I have is to load it up in Python and then save it as protobuf.

from tensorflow.python.tools import inspect_checkpoint as ckpt
ckpt.print_tensors_in_checkpoint_file('./resnet50/resnet_v2_50.ckpt',
                                      tensor_name='',
                                      all_tensors=False,
                                      all_tensor_names=False)

tf.reset_default_graph()
saver = tf.train.Saver()
with tf.Session() as sess:
    saver.restore(sess, './resnet50/resnet_v2_50.ckpt')
    print 'Model is restored'

    print 'Time to save the graph!'
    builder = tf.saved_model.builder.SavedModelBuilder('resnet_50_model')
    builder.add_meta_graph_and_variables(sess, ['serve'])
    builder.save()

However, this gives me an error saying that ValueError: No variables to save. I can fix it by declaring a variable

v1 = tf.get_variable('total_loss/ExponentialMovingAverage', shape=[])

But here comes my question, does that mean I must declare EVERY variable in the ResNet50 and have tensorflow load the values from ckpt file into those variables and then I perform save? Is there a shortcut to do this?

Answer

Peter Szoldan picture Peter Szoldan · Apr 1, 2018

You should import the variables as well from the .meta file if you have it available with the pretrained model like

saver = tf.train.import_meta_graph('./resnet50/resnet_v2_50.meta')

Detailed tutorial here.

If you don't have the .meta available, but you do have the network generating code, like with resnet_v2_50 in tensorflow/models/blob/master/research/slim/nets/resnet_v2.py then you should import that file and run the resnet_v2_50 function, which will define all the variables for you. Then restore the checkpoint.