Unable to restore a layer of class TextVectorization - Text Classification

user_8938979 0
Viewed 1788

System information Google Colab

When I run the example provided by official tensorflow basic text classification, everything runs fine until the model save, but when I load the model it gives me this error.

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.

Expected Behavior: Model should be loaded successfully and process the raw input

https://colab.research.google.com/gist/amahendrakar/8b65a688dc87ce9ca07ffb0ce50b84c7/44199.ipynb#scrollTo=fEjmSrKIqiiM

Example Link: https://tensorflow.google.cn/tutorials/keras/text_classification

text-classification tensorflow2.0 keras tensorflow python
2 Answers

I also ran into this error message (RuntimeError: Unable to restore a layer of class TextVectorization. [...]) when I implemented (and customized) the code from the "Basic Text Classification" tutorial.

Instead of running the code in a notebook, I have two scripts, one for building, training and saving the model and the other one for loading it and making predictions. (Thus, the error does not seem to be limited to Google Colab).

This is what I had to do (see https://github.com/tensorflow/tensorflow/issues/45231):

First, I added this line in the first script before the function definition and built, trained and saved the model again:

@tf.keras.utils.register_keras_serializable()
def custom_standardization(input_data):
[...]

# Save model as SavedModel
export_model.save(model_path, save_format='tf')

Secondly, I also had to add the same line and the whole function definition in the second script to make sure that it works if I restart(!) ipython (where I currently run the scripts) and only run the second script:

@tf.keras.utils.register_keras_serializable()
def custom_standardization(input_data):
      lowercase = tf.strings.lower(input_data)
      stripped_html = tf.strings.regex_replace(lowercase, '<br />', ' ')
      return tf.strings.regex_replace(stripped_html,
                                    '[%s]' % re.escape(string.punctuation),
                                    '')
[...]
# Load model
reloaded_model = tf.keras.models.load_model(model_path)
# Make predictions
predictions = reloaded_model.predict(examples)

Note: If I run the second script without restarting ipython after running the first script, I get this error:

ValueError: Custom>custom_standardization has already been registered [...]

Alternatively, you can just use the default standardization method in the vectorizer layer when building the model:

vectorize_layer = TextVectorization(
        standardize="lower_and_strip_punctuation",
        max_tokens=max_features,
        output_mode='int',
        output_sequence_length=sequence_length)
user_16870870
user_147581940
Related