CODEWITHSUNDEEP
pythontensorflow
debug_all_the_time
debug_all_the_time

Reputation: 574

How to save model in .pb format and then load it for inference in Tensorflow?

I am newbie of Tensorflow and trying to run one tutorial code located in https://github.com/Hvass-Labs/TensorFlow-Tutorials/blob/master/02_Convolutional_Neural_Network.ipynb

Based on this code, I would like to try to save the model in .pb format using simple_save and restore it for testing but I have no idea how to modify this piece of code. I have browsed some web pages but still didn't get the idea. Can anyone help me change this piece of code so that I can save the trained model and then load it for inference? Thank you!

Upvotes: 1

Views: 2625

Answers (1)

Mayur Bhangale
Mayur Bhangale

Reputation: 396

For saving model, you need two things - input and output tensor names. In your case, input tensor is called x and output tensor is y_pred and y_pred_cls (mentioned in In [29] in notebook). Here's simple example to save your model:

simple_save(session,
            export_dir,
            inputs={"x": x,},
            outputs={"y_pred": y_pred,
                     "y_pred_cls": y_pred_class})

EDIT: Restoring-

restoring_graph = tf.Graph()
with restoring_graph.as_default():
    with tf.Session(graph=restoring_graph) as sess:
       # Restore saved values
       tf.saved_model.loader.load(
          sess,
          [tag_constants.TRAINING],
          export_dir  # Path to SavedModel
       )
      # Pass inputs to model and do predictions below

Upvotes: 1

Related Questions