[BUG] Error when loading the session-based model in example

[FredHJC]

❓ Questions & Help

Details

https://github.com/NVIDIA-Merlin/models/blob/main/examples/usecases/ecommerce-session-based-next-item-prediction-for-fashion.ipynb

We are trying to implement the session-based models shown in the above notebook. However, there is a consistent error when loading the saved model. TypeError: ('Keyword argument not understood:', 'layer was saved without config')

We can just run the example notebook and save the trained model. Then the error occurs when trying to load it. It looks like customized layers / prediction tasks should be handled with a manually specified config.

[zhiruiwang]

Steps/Code to reproduce bug

1. Run the Dressipi notebook Transformer-based model example after the model is trained:

model_transformer.fit(loader, 
                      validation_data=val_loader,
                      epochs=EPOCHS
                     )

2. Save the model to disk, which is successful:

model_transformer.save(os.path.join('/ecom_modeling/11-30-session', 'transformer'))

It does have some warnings though:

WARNING:tensorflow:Skipping full serialization of Keras layer TFSharedEmbeddings(
  (_feature_shapes): Dict(
    (f_47_list_seq): TensorShape([1024, None])
    (f_68_list_seq): TensorShape([1024, None])
    (item_id_list_seq): TensorShape([1024, None])
    (item_id_last): TensorShape([1024, 1])
  )
  (_feature_dtypes): Dict(
    (f_47_list_seq): tf.int32
    (f_68_list_seq): tf.int32
    (item_id_list_seq): tf.int32
    (item_id_last): tf.int32
  )
), because it is not built.
.......
INFO:tensorflow:Unsupported signature for serialization: ((Prediction(outputs={'purchase_id_first/categorical_output': TensorSpec(shape=(None, 23272), dtype=tf.float32, name='outputs/outputs/purchase_id_first/categorical_output')}, targets={'purchase_id_first/categorical_output': TensorSpec(shape=(None, 23272), dtype=tf.float32, name='outputs/targets/purchase_id_first/categorical_output')}, sample_weight={'purchase_id_first/categorical_output': None}, features=None, negative_candidate_ids=None), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7fb613ff95e0>), {}).

3. Load the model from disk, which causes the error:

model_loaded = tf.keras.models.load_model(
        os.path.join('/ecom_modeling/11-30-session', 'transformer'))

TypeError: ('Keyword argument not understood:', 'layer was saved without config')

We are wondering if we did something wrong with the saving and loading of the model, or if there's a bug in saving and loading Merlin session-based models.

Also not sure if it's related to #898 or #889

[sararb]

Thank you for reporting the bug. 

I was able to reproduce this error when loading the model without importing merlin.models.tf.  In fact, the custom layers will not be understood by Tensorflow if Merlin Models (MM) is not imported. 

The recommended way to load a trained MM model is to import merlin.models.tf before calling tf.keras.models.load_model, as follows:

import tensorflow as tf
import merlin.models.tf as mm
model = tf.keras.models.load_model('transformer-model')

Please let us know if this fixes the loading issue. Thanks!

[oliverholworthy]

We also have a classmethod on our Model class. So can also do the following to load the model

import merlin.models.tf as mm

mm.Model.load('<path-to-saved-model-directory>')

[rnyak]

@zhiruiwang @FredHJC closing this issue since it should be solved via #927.

[zhiruiwang]

@rnyak We used the Merlin-tensorflow 22.12 image and refactored our pipeline to use the lastest API of merlin-models codebase, now the saving and loading two tower, LSTM, and transformer models are all working. Thanks for the help!