tensorflow2.0 - 使用 XLA tf.function 运行 TensorFlow 会引发错误

Question

当我尝试编译此代码时，出现以下错误。

File "xla_test.py", line 25, in <module>
    @tf.function(jit_compile=True)
TypeError: function() got an unexpected keyword argument 'jit_compile'

Answer 1

无需切换到 tf-nightly，只需使用：

@tf.function(experimental_compile=True)

从张量流文档：

Experimental_compile 如果为 True，则函数始终由 XLA 编译。XLA 在某些情况下可能更有效（例如 TPU、XLA_GPU、密集张量计算）。

在我的情况下，MCMC 采样在没有该参数的情况下进行：~1 分 37 秒，experimental_compile=True：~6 秒。从源代码构建的 TensorFlow（r2.4 分支）。

Answer 2

安装 tf-nightly 解决了这个问题。

pip install tf-nightly