26

官方@tf.function教程说：

要获得最佳性能并使您的模型可在任何地方部署，请使用 tf.function 从您的程序中制作图表。多亏了 AutoGraph，大量的 Python 代码只适用于 tf.function，但仍有一些陷阱需要提防。

主要内容和建议如下：

  • 不要依赖 Python 的副作用，如对象突变或列表追加。
  • tf.function 最适用于 TensorFlow 操作，而不是 NumPy 操作或 Python 原语。
  • 如有疑问，请使用 for x in y 成语。

它只提到了如何实现带@tf.function注释的函数，但没有提到何时使用它。

是否有关于如何决定我是否应该至少尝试用 注释函数的启发式方法tf.function？似乎没有理由不这样做，除非我懒得消除副作用或改变一些东西，比如range()-> tf.range()。但如果我愿意这样做...

有什么理由不使用@tf.function所有功能吗？

4

3 回答 3

29

TLDR：这取决于您的功能以及您是处于生产还是开发阶段。tf.function如果您希望能够轻松调试您的函数，或者如果它受到 AutoGraph 或 tf.v1 代码兼容性的限制，请不要使用。我强烈建议观看 Inside TensorFlow 关于AutoGraph和Functions 的讨论，而不是 Sessions。

下面我将分解原因，这些都是从谷歌在线提供的信息中获取的。

通常，tf.function装饰器会导致函数被编译为执行 TensorFlow 图的可调用对象。这需要：

  • 如果需要，通过 AutoGraph 转换代码（包括从注释函数调用的任何函数）
  • 跟踪并执行生成的图形代码

有关于这背后的设计理念的详细信息。

装饰函数的好处tf.function

一般福利

  • 更快的执行，特别是如果函数包含许多小操作（来源）

对于带有 Python 代码的函数/通过tf.function装饰使用 AutoGraph

如果你想使用 AutoGraph，tf.function强烈推荐使用而不是直接调用 AutoGraph。原因包括：自动控制依赖，某些 API 需要它，更多缓存和异常助手（来源）。

装饰函数的缺点tf.function

一般缺点

  • 如果该函数仅包含少量昂贵的操作，则不会有太多的加速（来源）

对于带有 Python 代码的函数/通过tf.function装饰使用 AutoGraph

  • 没有异常捕获（应该在急切模式下完成；在装饰函数之外）（来源）
  • 调试要困难得多
  • 由于隐藏的副作用和 TF 控制流造成的限制

提供有关 AutoGraph 限制的详细信息。

对于具有 tf.v1 代码的函数

  • 不允许在 中多次创建变量tf.function，但这可能会随着 tf.v1 代码的逐步淘汰而发生变化（来源）

对于带有 tf.v2 代码的函数

  • 没有特别的缺点

限制示例

多次创建变量

不允许多次创建变量，例如v以下示例：

@tf.function
def f(x):
    v = tf.Variable(1)
    return tf.add(x, v)

f(tf.constant(2))

# => ValueError: tf.function-decorated function tried to create variables on non-first call.

在以下代码中，通过确保self.v仅创建一次来缓解这种情况：

class C(object):
    def __init__(self):
        self.v = None
    @tf.function
    def f(self, x):
        if self.v is None:
            self.v = tf.Variable(1)
        return tf.add(x, self.v)

c = C()
print(c.f(tf.constant(2)))

# => tf.Tensor(3, shape=(), dtype=int32)

AutoGraph 未捕捉到的隐藏副作用

self.a无法隐藏此示例中的更改，这会导致错误，因为尚未完成跨功能分析（来源）：

class C(object):
    def change_state(self):
        self.a += 1

    @tf.function
    def f(self):
        self.a = tf.constant(0)
        if tf.constant(True):
            self.change_state() # Mutation of self.a is hidden
        tf.print(self.a)

x = C()
x.f()

# => InaccessibleTensorError: The tensor 'Tensor("add:0", shape=(), dtype=int32)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=cond_true_5, id=5477800528); accessed from: FuncGraph(name=f, id=5476093776).

显而易见的变化是没有问题的：

class C(object):
    @tf.function
    def f(self):
        self.a = tf.constant(0)
        if tf.constant(True):
            self.a += 1 # Mutation of self.a is in plain sight
        tf.print(self.a)

x = C()
x.f()

# => 1

TF 控制流的限制示例

这个 if 语句会导致错误，因为需要为 TF 控制流定义 else 的值：

@tf.function
def f(a, b):
    if tf.greater(a, b):
        return tf.constant(1)

# If a <= b would return None
x = f(tf.constant(3), tf.constant(2))   

# => ValueError: A value must also be returned from the else branch. If a value is returned from one branch of a conditional a value must be returned from all branches.
于 2020-04-30T05:30:08.760 回答
3

tf.function 在创建和使用计算图时很有用，它们应该在训练和部署中使用，但是大多数函数都不需要它。

假设我们正在构建一个特殊层，该层将与更大的模型分开。我们不希望在构造该层的函数之上放置 tf.function 装饰器，因为它只是对该层外观的定义。

另一方面，假设我们将进行预测或使用某些函数继续我们的训练。我们想要装饰器 tf.function 因为我们实际上是在使用计算图来获得一些价值。

一个很好的例子是构建一个编码器-解码器模型。不要将装饰器放在创建编码器或解码器或任何层的功能周围，这只是对它将做什么的定义。务必将装饰器放在“训练”或“预测”方法周围，因为它们实际上将使用计算图进行计算。

于 2020-04-30T05:07:31.763 回答
2

根据我的理解和文档，tf.function强烈建议使用主要是为了加速您的代码，因为被包装的代码tf.function将被转换为图形，因此有一些优化空间（例如操作修剪、折叠等）以当急切地运行相同的代码时，可能无法执行。

但是，在某些情况下，使用tf.function可能会产生额外的开销或不会导致明显的加速。一个值得注意的情况是包装函数很小并且只在代码中使用了几次，因此调用图形的开销可能相对较大。另一种情况是大多数计算已经在加速器设备（例如 GPU、TPU）上完成，因此图形计算获得的加速可能并不显着。

文档中还有一节讨论了各种场景下的加速，在本节的开头已经提到了上述两种情况：

仅仅包装一个使用张量的函数tf.function并不会自动加速你的代码。对于在单台机器上调用几次的小函数，调用图或图片段的开销可能会支配运行时。此外，如果大部分计算已经在加速器上进行，例如 GPU 密集型卷积堆栈，则图形加速不会很大。

对于复杂的计算，图可以提供显着的加速。这是因为图减少了 Python 到设备的通信并执行了一些加速。

但归根结底，如果它适用于您的工作流程，我认为为您的特定用例和环境确定这一点的最佳方法是分析您的代码何时以急切模式执行（即不使用tf.function）与何时执行它以图形模式执行（即tf.function广泛使用）。

于 2020-09-26T10:50:56.103 回答