[BUG] SavedModel serving signature contains all input features even when they aren't used

[karlhigley]

Bug description

When I train a model and save it, the resulting saved model file's serving signature contains every feature from the training dataset that was loaded by the dataloader, even when the model doesn't actually use those features.

Steps/Code to reproduce bug

Run the following code:

import tensorflow
import merlin.models.tf as mm
from merlin.datasets.synthetic import generate_data
from merlin.io import Dataset
from merlin.schema import Tags


train = generate_data("sequence-testing", num_rows=100)

seq_schema = train.schema.select_by_tag(Tags.SEQUENCE).select_by_tag(Tags.CATEGORICAL)
target = train.schema.select_by_tag(Tags.ITEM_ID).column_names[0]

predict_last = mm.SequencePredictLast(schema=seq_schema, target=target)

input_schema = seq_schema
output_schema = seq_schema.select_by_name(target)

loader = mm.Loader(train, batch_size=16, shuffle=False)

d_model = 48
query_encoder = mm.Encoder(
    mm.InputBlockV2(
        input_schema,
        embeddings=mm.Embeddings(
            input_schema.select_by_tag(Tags.CATEGORICAL), sequence_combiner=None
        ),
    ),
    mm.MLPBlock([d_model]),
    mm.GPT2Block(d_model=d_model, n_head=2, n_layer=2),
    tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=1)),
)

model = mm.RetrievalModelV2(
    query=query_encoder,
    output=mm.ContrastiveOutput(output_schema, negative_samplers="in-batch"),
)

model.compile(metrics={})
model.fit(loader, epochs=1, pre=predict_last)

query_encoder.save(
    f"/tmp/query_encoder"
)

Examine the serving signature with:

saved_model_cli show --tag_set serve --signature_def serving_default --dir /tmp/query_encoder

Expected behavior

The serving signature should only contain the features that the model actually uses. (Otherwise, when we try to serve the model, we have to provide a bunch of features that don't do anything.)

Environment details

• Merlin versions:
  • merlin-core 0.8.0+12.g8612b749e3
  • merlin-dataloader 0.0.2
  • merlin-models 0.9.0+42.g04597b9277
  • merlin-systems 0.7.0+19.g032be4d9
• Platform: Docker (merlin_ci_runner image)
• Python version: 3.8.10
• Tensorflow version (GPU?): tensorflow-gpu 2.9.2

Additional context

This is one of several issues that are currently making it difficult to serve Tensorflow session-based models for NVIDIA-Merlin/Merlin#433.

[rnyak]

this PR #904 is potentially solving this issue.