How to serve model in Tensorflow Serving that has custom function registered with @tf.keras.utils.register_keras_serializable?

user_774907 0
Viewed 906

I'm making use of the keras preprocessing layers:

@tf.keras.utils.register_keras_serializable(package='Custom', name=None)
def transform_domain(inputs):
    """ Lowercase domain and remove possible leading `www.`"""
    inputs = tf.strings.lower(inputs)
    return tf.strings.regex_replace(inputs, '^www\.', '', replace_global=False)

pre_layer_domain = TextVectorization(
    standardize=transform_domain, ...)

# The above layer is added to my model somewhere
# ...

model.fit(...)
model.save('out_path')

All is fine up to this point. The issue is when I try to load the model, I get an exception if my custom function is not present:

# In a new Python interpreter
model = tf.keras.models.load_model('out_path')

>>> RuntimeError:
    Unable to restore a layer of class TextVectorization. Layers of class
    TextVectorization require that the class be provided to the model loading
    code, either by registering the class using @keras.utils.register_keras_serializable
    on the class def and including that file in your program, or by passing the
    class in a keras.utils.CustomObjectScope that wraps this load call.

So the message suggests two things:

  • Use @keras.util.register_keras_serializable on the function while training the model + have that function loaded in the runtime while loading the model
  • Use the keras.utils.CustomObjectScope context manager when loading the model

Great, both options are lovely and wonderful - provided that I'm in control of how the model is loaded. But how do I do this when the model is loaded in Tensorflow Serving?

I've tried including a Python module with that function a few different ways by adding it to the out_path/assets directory. Can't figure out how to export that function so it'll somehow be loaded automatically when I load the model.

tensorflow-serving keras tensorflow
1 Answers

After some experimenting, I'm in fact able to export the model using either method:

model.save('out_path/dummy_model/1')
# or
tf.saved_model.save(model, 'out_path/dummy_model/1')

and when I load that model into Tensorflow Serving, "it just works." My custom function is loaded as part of the graph with no further work needed:

docker run --rm -p 8501:8501 \
  --mount type=bind,source=$(pwd)/out_path/,target=/app \
  -e MODEL_BASE_PATH=/app \
  -e MODEL_NAME=dummy_model -t tensorflow/serving

Inference on the appropriate endpoint for the example above (http://localhost:8501/v1/models/dummy_model:predict) works and returns the same values that I get when querying the model back in a Python environment, even though in Python land, I need to load the model using one of the two suggestions provided by the original error message I posted above (which in my case is not a problem, since when I load the model in Python, I have full control of the environment and the code that loads the model - typically a Notebook or other scripts for inspecting the internals of the model):

with tf.keras.utils.CustomObjectScope({'transform_domain': transform_domain}):
    new_model = tf.keras.models.load_model('dummy_model/1')
user_7749070
Related