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

Question

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？

Here is a minal 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():
    census = pd.read_csv("/Users/vincent/Projects/DeePray/examples/census/data/raw_data/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)

How can i use the reloaded tf.Keras.subclassing model within tf.feature_column layer to do inference？This puzzled me for days, Can anyone post an example to do that? I will be appreciated

Answer 1

This problem has beed solved by https://github.com/tensorflow/tensorflow/issues/43605#issuecomment-739076002