用于列表输入的Tensorflow 2 tf.function input_signature

Question

为了用saved_model api导出我的模型，我需要定义加载后要调用的每个方法的input_signature。我不知道如何判断输入是一个可变长度的列表(比如tf.keras.Model.call )。

下面列出了关于input_signature的一些未回答的问题：

• 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

还有关于*args：TensorFlow 2 How to use *args in tf.function?的这个，但是它不处理saved_model的问题。

Answer 1

也许你可以用张量代替列表作为输入？

然后在[None]中指定tf.TensorSpec维度，以允许跟踪重用方面的灵活性。

由于TensorFlow根据其形状匹配张量，因此使用None维度作为通配符将允许函数为可变大小的输入重用跟踪。如果您有不同长度的序列，或者对于每个批处理具有不同大小的图像，则可以发生可变大小的输入。

@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)