Save and load Universal Sentence Encoder model on different machines

How can we save and load an Universal Sentence Encoder model on different machines?

I created a Keras model with USE and saved it on machine A.

from tensorflow.keras.layers import Dense, Dropout, Input
from tensorflow.keras.models import Model, load_model
import tensorflow_hub as hub
import tensorflow as tf

module_url = "/path/on/machine/A/universal-sentence-encoder_4"
emb = hub.KerasLayer(module_url, input_shape=[], dtype=tf.string, trainable=True)

input1 = Input(shape=[], dtype=tf.string)
embedding_layer = emb(input1)
dense1 = Dense(units=512, activation="relu")(embedding_layer)
outputs = Dense(1, activation="sigmoid")(dense1)

model = Model(inputs=input1, outputs=outputs)
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["AUC"])
model.save("model.h5", include_optimizer=False)

Now I want to open model.h5 on machine B. Pre-trained USE is saved here /different/path/on/machine/B/universal-sentence-encoder_4. This is the error I get.

model = load_model("model.h5", custom_objects={"KerasLayer": hub.KerasLayer})

~/anaconda3/envs/tensorflow/lib/python3.8/site-packages/tensorflow_hub/resolver.py in __call__(self, handle)
    494   def __call__(self, handle):
    495     if not tf.compat.v1.gfile.Exists(handle):
--> 496       raise IOError("%s does not exist." % handle)
    497     return handle
    498 

OSError: /path/on/machine/A/universal-sentence-encoder_4 does not exist.

How can I resolve this issue? Is there a way to save everything including universal-sentence-encoder_4 into one model.h5 file so that users do not need to worry about USE?

tensorflow version: 2.4.1 keras version: 2.4.0

UPDATE: per WGierke's advice, created Google Colab to demonstrate the issue.

Solution

The issue here is related to what you've observed in this SO question as well: saving a Keras model requires serializing everything that is contained in the model. When you initialize a hub.KerasLayer using a generic Callable like loaded_obj, it cannot be serialized. Instead, you have to pass a string handle that points to the path of the SavedModel (or the tfhub.dev URL). When the Keras model is saved, KerasLayer.get_config is invoked, which stores that string in the config entry with the key handle.

When restoring the Keras model, the config is de-serialized and Keras runs hub.KerasLayer(config["handle"]). As you can see, this will fail if the model stored at handle is not available anymore.

Unfortunately, the only workaround right now is to make sure that the referenced path is also available on machine B.