Skip to content

Instantly share code, notes, and snippets.

@Mistobaan
Created June 6, 2020 05:06
Show Gist options
  • Save Mistobaan/dab7d05deccedf7670a4e1bd5a35c0ef to your computer and use it in GitHub Desktop.
Save Mistobaan/dab7d05deccedf7670a4e1bd5a35c0ef to your computer and use it in GitHub Desktop.
How to add a `key` parameter to serving models signatures in tensorflow
# Serving function that passes through keys
@tf.function(input_signature=[{
'is_male': tf.TensorSpec([None,], dtype=tf.string, name='is_male'),
'mother_age': tf.TensorSpec([None,], dtype=tf.float32, name='mother_age'),
'plurality': tf.TensorSpec([None,], dtype=tf.string, name='plurality'),
'gestation_weeks': tf.TensorSpec([None,], dtype=tf.float32, name='gestation_weeks'),
'key': tf.TensorSpec([None,], dtype=tf.string, name='key')
}])
def my_serve(inputs):
feats = inputs.copy()
key = feats.pop('key')
output = model(feats)
return {'key': key, 'babyweight': output}
tf.saved_model.save(model, EXPORT_PATH,
signatures={'serving_default': my_serve})
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment