使用@tffunction 的 Tensorflow2 警告

Question

此示例代码来自 Tensorflow 2

writer = tf.summary.create_file_writer("/tmp/mylogs/tf_function")

@tf.function
def my_func(step):
  with writer.as_default():
    # other model code would go here
    tf.summary.scalar("my_metric", 0.5, step=step)

for step in range(100):
  my_func(step)
  writer.flush()

但它正在发出警告。

警告：张量流：最近 5 次调用中的 5 次触发了 tf.function 回溯。跟踪是昂贵的，过多的跟踪可能是由于传递了 python 对象而不是张量。此外， tf.function 有 Experiment_relax_shapes=True 选项，可以放宽参数形状，以避免不必要的回溯。有关更多详细信息，请参阅 https://www.tensorflow.org/beta/tutorials/eager/tf_function#python_or_tensor_args 和https://www.tensorflow.org/api_docs/python/tf/function。

有一个更好的方法吗？

Answer 1

tf.function有一些“特殊性”。我强烈推荐阅读这篇文章：https : //www.tensorflow.org/tutorials/customization/performance

在这种情况下，问题在于每次您使用不同的输入签名调用该函数时都会“回溯”（即构建一个新图）。对于张量，输入签名指的是 shape 和 dtype，但对于 Python 数字，每个新值都被解释为“不同”。在这种情况下，因为您使用step每次都会更改的变量调用函数，所以每次都会回溯该函数。对于“真实”代码（例如在函数内部调用模型），这将非常慢。

您可以通过简单地转换step为张量来修复它，在这种情况下，不同的值不会算作新的输入签名：

for step in range(100):
    step = tf.convert_to_tensor(step, dtype=tf.int64)
    my_func(step)
    writer.flush()

或用于tf.range直接获取张量：

for step in tf.range(100):
    step = tf.cast(step, tf.int64)
    my_func(step)
    writer.flush()

这不应该产生警告（并且要快得多）。

Answer 2

尝试这个

我用model(x)而不是model.predict(x)它对我有用

如果您在自定义函数中遇到此错误，请为您的函数添加shape和的固定签名dtype。

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
...