TF2.3 load subclassing model within tf.feature_column layer get ValueError: Could not find matching function to call loaded from the SavedModel

[fuhailin]

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS Catalina
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: None
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): v2.3.0-54-gfcc4b966f1 2.3.1
• Python version: Python 3.7.4
• Bazel version (if compiling from source): None
• GCC/Compiler version (if compiling from source): None
• CUDA/cuDNN version: None
• GPU model and memory: None

You can collect some of this information using our environment capture script You can also obtain the TensorFlow version with:

1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"
2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior

I am trying to create a custom classification model using Tensorflow2.3 through tf.keras.Model subclassing method, in the subclass model init function, i use tf.feature_column layer to precess features. going through all of above part, i can train, save and reload the Saved_model, but when i use the reload model to do inference, i get the following error:

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (3 total):
    * {'age': [35], 'education': ['Bachelors']}
    * False
    * None
  Keyword arguments: {}

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

Option 1:
  Positional arguments (3 total):
    * {'education': TensorSpec(shape=(None, 1), dtype=tf.string, name='education'), 'age': TensorSpec(shape=(None, 1), dtype=tf.int64, name='age')}
    * True
    * None
  Keyword arguments: {}

Option 2:
  Positional arguments (3 total):
    * {'age': TensorSpec(shape=(None, 1), dtype=tf.int64, name='age'), 'education': TensorSpec(shape=(None, 1), dtype=tf.string, name='education')}
    * False
    * None
  Keyword arguments: {}

Option 3:
  Positional arguments (3 total):
    * {'age': TensorSpec(shape=(None, 1), dtype=tf.int64, name='inputs/age'), 'education': TensorSpec(shape=(None, 1), dtype=tf.string, name='inputs/education')}
    * False
    * None
  Keyword arguments: {}

Option 4:
  Positional arguments (3 total):
    * {'education': TensorSpec(shape=(None, 1), dtype=tf.string, name='inputs/education'), 'age': TensorSpec(shape=(None, 1), dtype=tf.int64, name='inputs/age')}
    * True
    * None
  Keyword arguments: {}

When I try to create model with tf.Keras.sequential class or without tf.feature_column layer, every thing works fine, so how can I use the reloaded tf.Keras.subclassing model within tf.feature_column layer to do inference？ This puzzled me for days.

Describe the expected behavior

tf.Keras.subclassing model within tf.feature_column layer can do inference like sequential model.

Standalone code to reproduce the issue Provide a reproducible test case that is the bare minimum necessary to generate the problem. If possible, please share a link to Colab/Jupyter/any notebook.

Jupyer minimum demo: https://www.kaggle.com/hailinfufu/notebookadf7121b80

Here is a minimal demo to reproduce my problem:

import pathlib
import time

import pandas as pd
import tensorflow as tf
from sklearn.model_selection import train_test_split

__SELECT_COLUMN_NAMES = ['age', 'education', 'income_bracket']

def get_train_test_pandas_data():
    # data can be download from: https://www.kaggle.com/uciml/adult-census-income?select=adult.csv
    census = pd.read_csv("adult_data.csv")

    census['income_bracket'] = census['income_bracket'].apply(lambda label: 0 if label == ' <=50K' else 1)
    census = census[__SELECT_COLUMN_NAMES]

    y_labels = census.pop('income_bracket')
    x_data = census

    x_train, x_test, y_train, y_test = train_test_split(x_data, y_labels, test_size=0.3)

    return x_train, x_test, y_train, y_test

def get_feature_columns():
    age = tf.feature_column.numeric_column("age", dtype=tf.int64)
    education = tf.feature_column.embedding_column(
        tf.feature_column.categorical_column_with_hash_bucket("education", hash_bucket_size=1000),
        dimension=100)

    feat_cols = [age, education]

    return feat_cols

if (tf.__version__ < '2.0'):
    tf.enable_eager_execution()

x_train, _, y_train, _ = get_train_test_pandas_data()

dataset = tf.data.Dataset.from_tensor_slices((dict(x_train), y_train))

dataset = dataset.shuffle(len(x_train)).batch(4)

feat_cols = get_feature_columns()

class mymodel(tf.keras.Model):
    def __init__(self):
        super(mymodel, self).__init__()
        self.layer1 = tf.keras.layers.DenseFeatures(feature_columns=feat_cols)
        self.layer2 = tf.keras.layers.Dense(10, activation='relu')
        self.layer3 = tf.keras.layers.Dense(10, activation='relu')
        self.layer4 = tf.keras.layers.Dense(1, activation='sigmoid')

    @tf.function
    def call(self, inputs, training=None, mask=None):
        x = self.layer1(inputs)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x

model = mymodel()

model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])

model.fit(dataset, epochs=1)

__SAVED_MODEL_DIR = './saved_models/census_keras/{}'.format(int(time.time()))
pathlib.Path(__SAVED_MODEL_DIR).mkdir(parents=True, exist_ok=True)

tf.saved_model.save(model, export_dir=__SAVED_MODEL_DIR)

you can replace model = mymodel() with

model = tf.keras.Sequential([
    tf.keras.layers.DenseFeatures(feature_columns=feat_cols),
    tf.keras.layers.Dense(10, activation='relu'),
    tf.keras.layers.Dense(10, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

that will work fine.

After trained and saved the model, i try to load the SavedModel to do predict:

import tensorflow as tf

# loaded_model = tf.keras.models.load_model("./saved_models/census_keras/1601196783")  # tf.saved_model.load("saved/1")
loaded_model = tf.keras.models.load_model("./saved_models/census_keras/1601196783")

y_pred = loaded_model.call({"age": [35],
                            "education": ["Bachelors"]})
print(y_pred)

y_pred = loaded_model.call({"age": [40],
                            "education": ["Assoc-voc"]})
print(y_pred)

Other info / logs Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

[Saduf2019]

@fuhailin Please provide with all dependencies to run the code shared or if possible share a colab gist with issue reported. I ran the code shared and face this error

[fuhailin]

@fuhailin Please provide with all dependencies to run the code shared or if possible share a colab gist with issue reported. I ran the code shared and face this error

Updated, please check this colab : https://colab.research.google.com/gist/fuhailin/32e3b834cff5856c53d719fe877414cb/untitled418.ipynb

[Saduf2019]

@fuhailin I cannot download the csv, please share the dataset for us to replicate the issue.

[fuhailin]

@fuhailin I cannot download the csv, please share the dataset for us to replicate the issue.

@Saduf2019 I have created a gist file to load, please check that colab again.

[Saduf2019]

I am able to replicate the issue reported, please find the gist here

[monicadsong]

Hi there--

You can do inference with your subclassed tf.keras.Model loaded_model by:

y_pred = loaded_model.call({"age": [[35]], "education": [["Bachelors"]]})
y_pred = loaded_model.call({"age": [[40]], "education": [["Assoc-voc"]]})

or

y_pred = loaded_model.call({"age": [[35], [40]],  "education": [["Bachelors"], ["Assoc-voc"]]})

The error message you received stated that it was expecting arguments "{'age': TensorSpec(shape=(None, 1), dtype=tf.int64, name='age'), 'education': TensorSpec(shape=(None, 1), dtype=tf.string, name='education')}", so the shape of your inputs should be of shape (None, 1).

You can check the shape by:

a = tf.constant([35])
b = tf.constant([[35]])
c = tf.constant([[35],[40]])
a.shape # TensorShape([1])
b.shape # TensorShape([1, 1])
c.shape  # TensorShape([2, 1])

The None in the TensorSpec shape parameter means that the batch size is variable.

Closing this issue, but please reopen if I overlooked something. Thanks.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No