为何TensorFlow无参数@tf.function设置固定随机种子无效？

测试代码 #

tf.random.set_seed(1234)

@tf.function
def g(i):
  print("round", i)
  return tf.random.uniform([1])

@tf.function
def f():
  return tf.random.uniform([1])

print(g(0))
print(g(1))
print(f())
print(f())

执行输出 #

round 0
tf.Tensor([0.13047123], shape=(1,), dtype=float32)
round 1
tf.Tensor([0.13047123], shape=(1,), dtype=float32)
tf.Tensor([0.5380393], shape=(1,), dtype=float32)
tf.Tensor([0.3253647], shape=(1,), dtype=float32)

原因解析 #

• 带参数的函数g：@tf.function会根据输入参数的类型或值决定是否重新追踪函数并构建新计算图。调用g(0)和g(1)时，传入的参数不同，会触发两次独立的图追踪过程。每次追踪都会基于你设置的全局随机种子1234初始化随机生成器状态，因此两次生成的随机值完全一致，两次打印"round"也证明函数被重新追踪了。
• 无参数的函数f：第一次调用f()时，@tf.function完成了图的追踪与构建，后续所有调用都会复用这个已构建好的计算图。在TensorFlow图模式中，tf.random.uniform这类随机操作的状态会持久化在图内，每次调用函数时随机生成器的状态会自动推进，所以每次生成的随机值都不一样。

内容的提问来源于stack exchange，提问作者zell