The model does not save and load correctly when containing tf.keras.layers.experimental.preprocessing.StringLookup layer

[oawxkw]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): v2.14.0-rc0-34-gdd01672d9a9 2.14.0-rc1
• Python version:
• Bazel version (if compiling from source):
• GPU model and memory:
• Exact command to reproduce:

Describe the problem.

The model does not save and load correctly when containing tf.keras.layers.experimental.preprocessing.StringLookup layer.
It seems that the vocabulary is not saved or loaded correctly, which is empty when loading the model.
I manually checked the saved model files and found that:
The vocabulary is saved as a constant value in saved_model.pb and all layers arguments are dumped in keras_metadata.pb, but it's not saved in variables/variables.data-00000-of-00001 and variables/variables.index (cannot find the sting aaaa or bbbb in these two files).
If the saved model files are correct, there should be something wrong when loading the model.

I've reported this issue in tensorflow/tensorflow#61779, but the contributor suggested reporting it here.
This behavior may also relate to tensorflow/tensorflow#61369, but in a different API endpoint.

Describe the current behavior.

Loading the saved model throws an error, which is caused by the empty vocabulary argument of the StringLookup layer.

Describe the expected behavior.

Saving and loading the model should work correctly for the StringLookup layer.

Standalone code to reproduce the issue.

import pickle
import tensorflow as tf
print(tf.version.GIT_VERSION, tf.version.VERSION, flush=True)

model_input = tf.keras.Input(shape=(1,), dtype=tf.int64)
lookup = tf.keras.layers.experimental.preprocessing.StringLookup(vocabulary=['aaaaaaaaaaaaaaaa', 'bbbbbbbbbbbbbbbb'])(model_input)
output = tf.keras.layers.Dense(10)(lookup)
full_model = tf.keras.Model(model_input, output)

# this part works
try:
    print(full_model.layers[1].get_config())
    model_bytes = pickle.dumps(full_model)
    model_recovered = pickle.loads(model_bytes)
except Exception as e:
    print("Failed! Error:", e, flush=True)
else:
    print("Success!", flush=True)

# this part throws an error
try:
    full_model.save("/tmp/temp_model")
    full_model_loaded = tf.keras.models.load_model("/tmp/temp_model")
    print(full_model_loaded.layers[1].get_config())
    model_bytes = pickle.dumps(full_model_loaded)
    model_recovered = pickle.loads(model_bytes)
except Exception as e:
    print("Failed! Error:", e, flush=True)
else:
    print("Success!", flush=True)

Source code / logs.

v2.14.0-rc0-34-gdd01672d9a9 2.14.0-rc1
{'name': 'string_lookup', 'trainable': True, 'dtype': 'int64', 'invert': False, 'max_tokens': None, 'num_oov_indices': 1, 'oov_token': '[UNK]', 'mask_token': None, 'output_mode': 'int', 'sparse': False, 'pad_to_max_tokens': False, 'idf_weights': None, 'vocabulary': ListWrapper(['aaaaaaaaaaaaaaaa', 'bbbbbbbbbbbbbbbb']), 'vocabulary_size': 3, 'encoding': 'utf-8'}
Success!
{'name': 'string_lookup', 'trainable': True, 'dtype': 'int64', 'invert': False, 'max_tokens': None, 'num_oov_indices': 1, 'oov_token': '[UNK]', 'mask_token': None, 'output_mode': 'int', 'sparse': False, 'pad_to_max_tokens': False, 'idf_weights': None, 'vocabulary': ListWrapper([]), 'vocabulary_size': 3, 'encoding': 'utf-8'}
Failed! Error: Error when deserializing class 'StringLookup' using config={'name': 'string_lookup', 'trainable': True, 'dtype': 'int64', 'invert': False, 'max_tokens': None, 'num_oov_indices': 1, 'oov_token': '[UNK]', 'mask_token': None, 'output_mode': 'int', 'sparse': False, 'pad_to_max_tokens': False, 'idf_weights': None, 'vocabulary': [], 'vocabulary_size': 3, 'encoding': 'utf-8'}.

Exception encountered: Cannot set an empty vocabulary, you passed [].

[nkovela1]

Hi @oawxkw, please try the new .keras saving format that is now the recommended default for saving and loading in Keras. Simply add .keras at the end of your filepath. I've checked and it solves the issue: https://colab.research.google.com/gist/nkovela1/db2e1463c84697e1f378278195e7fdf7/18376_bug_repro.ipynb

[oawxkw]

Hello @nkovela1, thanks for your reply.

I think this issue may relate to the format of the models, which is causing the vocabulary not saved or loaded correctly.
It seems that the keras format can deal with the vocabulary correctly, but the tf format cannot.
I've read the documentation of tf.keras.Model#save and tf.keras.saving.save_model and tried to apply the different save_format parameters with some filepath extensions, but it seems that the save_format parameter does not work as expected.

• filepath: str or pathlib.Path object. Path where to save the model.
• save_format: Either "keras", "tf", "h5", indicating whether to save the model in the native Keras format (.keras), in the TensorFlow SavedModel format (referred to as "SavedModel" below), or in the legacy HDF5 format (.h5). Defaults to "tf" in TF 2.X, and "h5" in TF 1.X.

According to the documentation, the save_format parameter should be able to force the model to be saved in the specified format, while the filepath extension should be ignored.
But, I found the source code of tf.keras.saving.save_model, which seems to be inconsistent with the documentation, where L130 only checks the save_format or filepath parameters of h5, but when passed to L147 the "save_format" parameter is ignored.

https://github.com/tensorflow/tensorflow/blob/8cf6ece80115da66eb7286a04c9722d6398c26e5/tensorflow/python/keras/saving/save.py#L120-L148

I'm using the following code to test the different save_format arguments with different filepath extensions.
Hope it can help to reproduce the issue.

import pickle
import tensorflow as tf
print(tf.version.GIT_VERSION, tf.version.VERSION, flush=True)


model_input = tf.keras.Input(shape=(1,), dtype=tf.int64)
lookup = tf.keras.layers.experimental.preprocessing.StringLookup(vocabulary=['aaaaaaaaaaaaaaaa', 'bbbbbbbbbbbbbbbb'])(model_input)
output = tf.keras.layers.Dense(10)(lookup)
full_model = tf.keras.Model(model_input, output)


def test_save(filepath, save_format=None):
    try:
        full_model.save(filepath, save_format=save_format)
        full_model_loaded = tf.keras.models.load_model(filepath)
        print(filepath, save_format, full_model_loaded.layers[1].get_config())
        model_bytes = pickle.dumps(full_model_loaded)
        model_recovered = pickle.loads(model_bytes)
    except Exception as e:
        print("Failed! Error:", e, flush=True)
    else:
        print("Success!", flush=True)


test_save("/tmp/temp_model", save_format=None)
test_save("/tmp/temp_model.keras", save_format=None)
test_save("/tmp/temp_model.tf", save_format=None)

test_save("/tmp/temp_model", save_format='keras')
test_save("/tmp/temp_model.keras", save_format='keras')
test_save("/tmp/temp_model.tf", save_format='keras')

test_save("/tmp/temp_model", save_format='tf')
test_save("/tmp/temp_model.keras", save_format='tf')
test_save("/tmp/temp_model.tf", save_format='tf')