为什么 Tensorflow 函数对函数的不同整数输入执行回溯？

Question

我正在关注函数here 上的 Tensorflow 指南，根据我的理解，TF 将为每次调用具有不同输入签名（即数据类型和输入形状）的函数创建一个图表。但是，以下示例使我感到困惑。由于两个输入都是整数并且具有完全相同的形状，因此 TF 不应该只执行一次跟踪和构造图吗？为什么调用函数时会发生两次跟踪？

@tf.function
def a_function_with_python_side_effect(x):
  print("Tracing!") # An eager-only side effect.
  return x * x + tf.constant(2)


# This retraces each time the Python argument changes,
# as a Python argument could be an epoch count or other
# hyperparameter.
print(a_function_with_python_side_effect(2))
print(a_function_with_python_side_effect(3))

输出：

Tracing!
tf.Tensor(6, shape=(), dtype=int32)
Tracing!
tf.Tensor(11, shape=(), dtype=int32)

Answer 1

数字 2 和 3 被视为不同的整数值，这就是您看到“正在跟踪！”的原因两次。您所指的行为：“TF 将为每次调用具有不同输入签名（即数据类型和输入形状）的函数创建一个图形”适用于张量而不是简单数字。您可以通过将两个数字都转换为张量常数来验证这一点：

import tensorflow as tf

@tf.function
def a_function_with_python_side_effect(x):
  print("Tracing!") # An eager-only side effect.
  return x * x + tf.constant(2)

print(a_function_with_python_side_effect(tf.constant(2)))
print(a_function_with_python_side_effect(tf.constant(3)))

Tracing!
tf.Tensor(6, shape=(), dtype=int32)
tf.Tensor(11, shape=(), dtype=int32)

这是混合 python 标量和tf.function 时的副作用。查看here的追踪规则。你读到了：

为 tf.Tensor 生成的 cache 键是它的形状和 dtype。

为 Python 原语（如 int、float、str）生成的 cache 键是它的值。