Tensorflow 2 tf.function input_signature for a list input

Question

In order to export my model with the saved_model api, I need to define the input_signature of each method intended to be called after loading. I don't know how to tell that the input is a list with variable length (as it is for tf.keras.Model.call for instance).

There is a list of unanswered questions about input_signature on SO:

• Use dictionary in tf.function input_signature in Tensorflow 2.0
• tf.function input_signature for distributed dataset in tensorflow 2.0
• TensorFlow 2 tf.function decorator

and also this one about *args : TensorFlow 2 How to use *args in tf.function? but it does not handle the problem of saved_model .

Answer 1

Maybe you could use a Tensor instead of a list as input?

Then specify a [None] dimension in tf.TensorSpec to allow for flexibility in trace reuse.

Since TensorFlow matches tensors based on their shape, using a None dimension as a wildcard will allow Functions to reuse traces for variably-sized input. Variably-sized input can occur if you have sequences of different length, or images of different sizes for each batch.

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def g(x):
  print('Tracing with', x)
  return x

# No retrace!
print(g(tf.constant([1, 2, 3])))
print(g(tf.constant([1, 2, 3, 4, 5])))

Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([1 2 3], shape=(3,), dtype=int32)
tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)