Search code examples
modeltensorflowpredictionmnisttf-slim

Tensorflow Slim restore model and predict


I'm currently trying to learn how to use TF-Slim and I'm following this tutorial: https://github.com/mnuke/tf-slim-mnist.

Assuming that I already have a trained model saved in a checkpoint, how do I now use that model and apply it? Like, in the tutorial how do I use my trained MNIST model and feed in a new set of MNIST images, and print the predictions?


Solution

  • You can try a workflow like:

    #obtain the checkpoint file
    checkpoint_file= tf.train.latest_checkpoint("./log")
    
    #Construct a model as such:
    with slim.arg_scope(mobilenet_arg_scope(weight_decay=weight_decay)):
                logits, end_points = mobilenet(images, num_classes = dataset.num_classes, is_training = True, width_multiplier=width_multiplier)
    
    #Obtain the trainable variables and a saver
    variables_to_restore = slim.get_variables_to_restore()
    saver = tf.train.Saver(variables_to_restore)
    
    #Proceed to create your training optimizer and predictions monitoring, summaries etc.
    ...
    
    #Finally when you're about to train your model in a session, restore the checkpoint model to your graph first:
    
    with tf.Session() as sess:
        saver.restore(sess, checkpoint_file)
        #...Continue your training
    

    Basically you have to get the right variables to be restored, and these variables must have names that match those found in your checkpoint model. Afterwards, pass the list of variables to be restored to a Saver, and then in a TF session, let the saver restore all the variables from a checkpoint model in the session.