Tensorflow 2 渴望与图形模式

数据挖掘 张量流
2021-10-12 00:41:53

我一直在学习 tensorflow-2.0.0 beta 教程。高级示例tensorflow.keras中,使用了一个子类。@tf.function装饰器train_step存在test_step意味着模型在graph模式下执行(不确定这是否是正确的术语,我的意思是与eager模式相反)。如果我删除这些装饰器,我可以直接进入模型call函数并查看每一层的输入/输出张量,这很整洁。

我的问题是,是否有一种编程方式来启用/禁用@tf.function装饰器。将它们注释掉以在急切模式和图形模式之间切换似乎并不是特别可扩展,但它对于调试/学习肯定很有用)

2个回答

您总是可以编写两个函数(一个带有装饰器,一个不带)并调用适合您的函数...

例如

@tf.function()
def graph_function()
    # This function will operate in graph mode
    ...

def eager_function()
    # This function will operate in eager mode
   ...

if tf.executing_eagerly()
    my_function = eager_function
else:
    my_function = graph_function

# You proceed to my_function from now on

我不知道是否有更好的方法,但我在 github 上的 tensorflow 官方存储库中看到了很多这种方法。

我不知道它是在哪个版本的 Tensorflow 中引入的,但至少在 TF 2.1 中,有

tf.config.experimental_run_functions_eagerly(True)

可用的。它使所有@tf.function-decorated 函数无论如何都以急切模式运行,直到通过False再次调用 with 参数将其重置。

有关详细信息,请参阅Tensorflow Experimental_run_functions_eagerly 文档