Keras model couldn't save

[yzhang5471]

platform: (CoLab) uname_result(system='Linux', node='1acc1ece6828', release='4.19.104+', version='#1 SMP Wed Feb 19 05:26:34 PST 2020', machine='x86_64', processor='x86_64')

python==3.6.9
tensorflow==2.1.0
tensorflow-ranking==0.3.0

TFRanking custom Keras layers can serialize and deserialize but the custom model failed saving to either saved_model or h5 format. Could there be any issues with get_config implementations?

_LABEL_FEATURE = "relevance"
_PADDING_LABEL = -1
_SIZE="example_list_size"

def create_feature_columns():
  sparse_column = tf.feature_column.categorical_column_with_hash_bucket(
      key="user_id", hash_bucket_size=100, dtype=tf.int64)
  query_embedding = tf.feature_column.embedding_column(
      categorical_column=sparse_column, dimension=20)
  context_feature_columns = {"user_id": query_embedding}

  sparse_column = tf.feature_column.categorical_column_with_hash_bucket(
      key="document_id", hash_bucket_size=100, dtype=tf.int64)
  document_embedding = tf.feature_column.embedding_column(
      categorical_column=sparse_column, dimension=20)
  example_feature_columns = {"document_id": document_embedding}

  return context_feature_columns, example_feature_columns

context_feature_columns, example_feature_columns = create_feature_columns()

# instantiate keras model
network = tfr.keras.canned.DNNRankingNetwork(
    context_feature_columns=context_feature_columns,
    example_feature_columns=example_feature_columns,
    hidden_layer_dims=[1024, 512, 256],
    activation=tf.nn.relu,
    dropout=0.5)
ranker = tfr.keras.model.create_keras_model(
    network=network,
    loss=tfr.keras.losses.get(tfr.losses.RankingLossKey.SOFTMAX_LOSS),
    metrics=tfr.keras.metrics.default_keras_metrics(),
    optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.05),
    size_feature_name=_SIZE)

# save keras model to saved_model format
ranker.save('tmp')

INFO:tensorflow:Assets written to: tmp/assets
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-90-c6e3a44bcc2b> in <module>()
----> 1 ranker.save('tmp')

11 frames
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/util/serialization.py in get_json_type(obj)
     70     return obj.__wrapped__
     71 
---> 72   raise TypeError('Not JSON Serializable:', obj)

TypeError: ('Not JSON Serializable:', tf.int64)

[ramakumar1729]

@yzhangswingman : In the next release, we will have a test case that shows how to create saved models from Keras.

[yzhang5471]

@ramakumar1729 Looking forward to it! Any timeline on the new release?

[ramakumar1729]

@yzhangswingman : Release v0.3.1 has a commit which has a test case to create saved model from Keras. PTAL.

[yzhang5471]

@ramakumar1729 Cool, thanks for the update!

In my experience, when I serialized keras.Model to SavedModel without providing explicit signatures, the exported serving_default works fine even with SparseTensor input. In model_test.py however, it seems that custom signatures don't work well with SaprseTensor input?

BTW, keras.Model recreated from SavedModel seems to be not working properly. Would you consider supporting and testing this use case?

[ramakumar1729]

@yzhangswingman : Can you share minimum reproducible code for the issue you are facing? Also, it would be great if you can share how to create SavedModels with sparse tensor inputs.

[yzhang5471]

@ramakumar1729 sure, I will get back to you later with some code and stuff

[AlexanderJLiu]

@ramakumar1729 I've tried to save tfr keras model follow this commit, but failed.

Errors are shown as below:

ValueError: If specifying TensorSpec names for nested structures, either zero or all names have to be specified.

The reproducible code in cloab is here

Thanks~

[ramakumar1729]

@yzhangswingman @AlexanderJLiu : Right, this issue that comes up as the features are sparse (query, document). That is why I added a TODO that the saved model works only for dense features currently.

[AlexanderJLiu]

@ramakumar1729 Thanks, I didn't notice the TODO comments before.

[yzhang5471]

@ramakumar1729 @AlexanderJLiu
I came across the same error with TF2.2.0, yet the model exported just fine with TF2.1.0

!pip install tensorflow-ranking
!pip install -U tensorflow==2.1.0

import tensorflow as tf
import tensorflow_ranking as tfr

assert tf.version.VERSION == '2.1.0'

_LABEL_FEATURE = "relevance"
_PADDING_LABEL = -1
_SIZE="example_list_size"


def create_feature_columns():
  sparse_column = tf.feature_column.categorical_column_with_hash_bucket(
      key="user_id", hash_bucket_size=100, dtype=tf.int64)
  query_embedding = tf.feature_column.embedding_column(
      categorical_column=sparse_column, dimension=20)
  context_feature_columns = {"user_id": query_embedding}

  sparse_column = tf.feature_column.categorical_column_with_hash_bucket(
      key="document_id", hash_bucket_size=100, dtype=tf.int64)
  document_embedding = tf.feature_column.embedding_column(
      categorical_column=sparse_column, dimension=20)
  example_feature_columns = {"document_id": document_embedding}

  return context_feature_columns, example_feature_columns

context_feature_columns, example_feature_columns = create_feature_columns()

# instantiate keras layers and model
network = tfr.keras.canned.DNNRankingNetwork(
    context_feature_columns=context_feature_columns,
    example_feature_columns=example_feature_columns,
    hidden_layer_dims=[1024, 512, 256],
    activation=tf.nn.relu,
    dropout=0.5)
ranker = tfr.keras.model.create_keras_model(
    network=network,
    loss=tfr.keras.losses.get(tfr.losses.RankingLossKey.SOFTMAX_LOSS),
    metrics=tfr.keras.metrics.default_keras_metrics(),
    optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.05),
    size_feature_name=_SIZE)
ranker.save('tmp/test1')

WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1786: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
INFO:tensorflow:Assets written to: tmp/test1/assets

imported_model = tf.saved_model.load('tmp/test1')
serving_default_fn = imported_model.signatures['serving_default']

The exported forward pass function however expects all inputs to be dense Tensor.

test_input = {
      "user_id":
          tf.convert_to_tensor(value=[[1], [2]], dtype=tf.int64),
      "document_id":
          tf.SparseTensor(
              indices=[[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0]],
              values=tf.convert_to_tensor([0, 1, 2, 3], dtype=tf.int64),
              dense_shape=[2, 2, 1]),
      "example_list_size": tf.convert_to_tensor([2, 2], dtype=tf.int32)
  }

serving_default_fn(**test_input)

ValueError: All inputs to `ConcreteFunction`s must be Tensors; on invocation of signature_wrapper, the 0-th input (SparseTensor(indices=tf.Tensor(
[[0 0 0]
 [0 1 0]
 [1 0 0]
 [1 1 0]], shape=(4, 3), dtype=int64), values=tf.Tensor([0 1 2 3], shape=(4,), dtype=int64), dense_shape=tf.Tensor([2 2 1], shape=(3,), dtype=int64))) was not a Tensor.

test_input_all_dense = test_input.copy()
test_input_all_dense['document_id'] = tf.sparse.to_dense(test_input_all_dense['document_id'])

serving_default_fn(**test_input_all_dense)

{'dnn_ranking_network': <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
 array([[ 0.01578074,  0.01597136],
        [-0.0111532 , -0.02785628]], dtype=float32)>}

[yzhang5471]

@ramakumar1729 Here's the other error in reconstructing keras.Model object from SavedModel.

loaded_keras_model = tf.keras.models.load_model('tmp/test1')

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (1 total):
    * {'user_id': <tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fa24f5b7160>, 'document_id': <tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x7fa24f615c18>, 'example_list_size': <tf.Tensor 'inputs_3:0' shape=(None,) dtype=int32>}
  Keyword arguments: {}

Expected these arguments to match one of the following 1 option(s):

Option 1:
  Positional arguments (1 total):
    * {'example_list_size': TensorSpec(shape=(None,), dtype=tf.int32, name='inputs/example_list_size'), 'user_id': TensorSpec(shape=(None, None), dtype=tf.int64, name='inputs/user_id'), 'document_id': TensorSpec(shape=(None, None, None), dtype=tf.int64, name='inputs/document_id')}
  Keyword arguments: {}

[TalhaAsmal]

@ramakumar1729 I've tried to save tfr keras model follow this commit, but failed.

Errors are shown as below:

ValueError: If specifying TensorSpec names for nested structures, either zero or all names have to be specified.

The reproducible code in cloab is here

Thanks~

This error seems to have been resolved in the latest Tensorflow build. I had the same issue, but installing tf-nightly and tf-estimator-nightly resolved it.

[ramakumar1729]

Thanks @TalhaAsmal ! @yzhangswingman : Does this resolve your issue?

[yzhang5471]

@ramakumar1729: yeah, TF2.1/2.3 + TFR0.3.1 are good enough for the moment. Thanks!
also appreciate your input! @TalhaAsmal