Tensorflow 2 Keras API SavedModel example

savedmodel-tensorflow-2.py

import tensorflow as tf

# Define Keras model

inputs = tf.keras.Input(shape=(3,))
x = tf.keras.layers.Dense(4, activation=tf.nn.relu)(inputs)
outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax, name='output_name')(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

# Train it
# ...

# Test it
x = tf.constant([[1.0,2.0,3.0], [1.0,2.0,3.0]])
model.predict(x) 
# out -> array([[0.22246964, 0.00502986, 0.578113  , 0.16929728, 0.02509018],
#       [0.22246964, 0.00502986, 0.578113  , 0.16929728, 0.02509018]],
#      dtype=float32)

# Save it 
tf.saved_model.save(model, 'model')

# Load and list signature defs
loaded = tf.saved_model.load('model')
list(loaded.signatures.keys())
# out -> ['serving_default']

# Test it using predict function 
import numpy as np
predict = loaded.signatures["serving_default"]
some_data = np.random.rand(2,3)*5
result = predict(tf.constant(some_data, dtype=tf.float32))
result['output_name'].numpy()
# out -> array([[0.2900614 , 0.00554972, 0.38926545, 0.27773553, 0.03738797],
#       [0.3551294 , 0.00153254, 0.3601567 , 0.25889066, 0.02429073]],
#      dtype=float32)