how to load and use a saved model on tensorflow?

user_2210667 0
Viewed 33491

I have found 2 ways to save a model in Tensorflow: tf.train.Saver() and SavedModelBuilder. However, I can't find documentation on using the model after it being loaded the second way.

Note: I want to use SavedModelBuilder way because I train the model in Python and will use it at serving time in another language (Go), and it seems that SavedModelBuilder is the only way in that case.

This works great with tf.train.Saver() (first way):

model = tf.add(W * x, b, name="finalnode")

# save
saver = tf.train.Saver()
saver.save(sess, "/tmp/model")

# load
saver.restore(sess, "/tmp/model")

# IMPORTANT PART: REALLY USING THE MODEL AFTER LOADING IT
# I CAN'T FIND AN EQUIVALENT OF THIS PART IN THE OTHER WAY.

model = graph.get_tensor_by_name("finalnode:0")
sess.run(model, {x: [5, 6, 7]})

tf.saved_model.builder.SavedModelBuilder() is defined in the Readme but after loading the model with tf.saved_model.loader.load(sess, [], export_dir)), I can't find documentation on getting back at the nodes (see "finalnode" in the code above)

tensorflow
4 Answers

Here's the code snippet to load and restore/predict models using the simple_save

#Save the model:
tf.saved_model.simple_save(sess, export_dir=saveModelPath,
                                   inputs={"inputImageBatch": X_train, "inputClassBatch": Y_train,
                                           "isTrainingBool": isTraining},
                                   outputs={"predictedClassBatch": predClass})

Note that using simple_save sets certain default values (this can be seen at: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/simple_save.py)

Now, to restore and use the inputs/outputs dict:

from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model import signature_constants

with tf.Session() as sess:
  model = tf.saved_model.loader.load(export_dir=saveModelPath, sess=sess, tags=[tag_constants.SERVING]) #Note the SERVINGS tag is put as default.

  inputImage_name = model.signature_def[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].inputs['inputImageBatch'].name
  inputImage = tf.get_default_graph().get_tensor_by_name(inputImage_name)

  inputLabel_name = model.signature_def[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].inputs['inputClassBatch'].name
  inputLabel = tf.get_default_graph().get_tensor_by_name(inputLabel_name)

  isTraining_name = model.signature_def[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].inputs['isTrainingBool'].name
  isTraining = tf.get_default_graph().get_tensor_by_name(isTraining_name)

  outputPrediction_name = model.signature_def[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].outputs['predictedClassBatch'].name
  outputPrediction = tf.get_default_graph().get_tensor_by_name(outputPrediction_name)

  outPred = sess.run(outputPrediction, feed_dict={inputImage:sampleImages, isTraining:False})

  print("predicted classes:", outPred)

Note: the default signature_def was needed to make use of the tensor names specified in the input & output dicts.

user_104463960

A code snippet that worked for me to load a pb file and inference on a single image.

The code follows the following steps: load the pb file into a GraphDef (a serialized version of a graph (used to read pb files), load GraphDef into a Graph, get input and output tensors by their names, inference on a single image.

import tensorflow as tf 
import numpy as np
import cv2

INPUT_TENSOR_NAME = 'input_tensor_name:0'
OUTPUT_TENSOR_NAME = 'output_tensor_name:0'

# Read image, get shape
# Add dimension to fit batch shape
img = cv2.imread(IMAGE_PATH)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = img.astype(float)
height, width, channels = image.shape
image = np.expand_dims(image, 0)  # Add dimension (to fit batch shape)


# Read pb file into the graph as GraphDef - Serialized version of a graph     (used to read pb files)
with tf.gfile.FastGFile(PB_PATH, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

# Load GraphDef into Graph
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def, name="")

# Get tensors (input and output) by name
input_tensor = graph.get_tensor_by_name(INPUT_TENSOR_NAME)
output_tensor = graph.get_tensor_by_name(OUTPUT_TENSOR_NAME)

# Inference on single image
with tf.Session(graph=graph) as sess:
    output_vals = sess.run(output_tensor, feed_dict={input_tensor: image})  #
user_129204410
Related