我应该将@tf.function 用于所有功能吗？

Posted 2023-02-16

【中文标题】我应该将@tf.function 用于所有功能吗？

【英文标题】：Should I use @tf.function for all functions?

【发布时间】：2020-05-07 20:45:37

【问题描述】：

@tf.function 上的 official tutorial 说：

为了获得最佳性能并使您的模型可在任何地方部署， 使用 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？

【问题讨论】：

为什么要添加这些标签？我们也可以添加tensorflow0.1、tensorflow0.2、tensorflow0.3、tensorflow0.4、tensorflow0.5 等等，然后为每个these tf modules and classes 添加标签。另外，为什么不为每个 Python 的标准模块及其函数和类添加一个标签呢？

这就是我引入tensorflow2.x标签的原因，因为有些问题不仅与tensorflow2.0有关，而且与tensorflow2.x标签有关。但是，为每个版本的库添加标签是不合适和不可行的。以 Python 为例。你没有python3.4.6.....python.3.8.2，但是python3.x

一方面，tf.function guide 说“装饰模块级函数和模块级类的方法，避免装饰局部函数或方法”。我似乎记得更明确的措辞，比如“不要装饰每个函数，在更高级别的函数中使用tf.function，比如训练循环”，但我可能记错了（或者它可能已被删除）。 OTOH，this discussion 有来自开发人员的有趣输入，最后似乎可以在任何张量/变量函数中使用它。

@jdehesa AFAIK @tf.function 带注释的函数还将它们自称为图形的函数编译。因此，您只需要注释与您描述的内容一致的模块的入口点。但是手动注释调用堆栈中较低的函数也没有什么坏处。

@problemofficer 是的，所以在我链接的 GitHub 问题中，有一些关于创建多个中间函数是否会对性能产生轻微影响的讨论，但似乎图形优化器（抓取器）可以“内联”如果需要，可以使用函数，但另一方面，如果多次调用另一个非tf.function，则无法防止图中的“代码重复”，这就是为什么建议广泛使用的原因。

【参考方案1】：

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

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

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

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

There is detailed information available on the design ideas behind this.

用tf.function装饰函数的好处

一般福利

更快的执行，尤其是当函数包含许多小操作时(Source)

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

如果您想使用 AutoGraph，强烈建议使用 tf.function，而不是直接调用 AutoGraph。 造成这种情况的原因包括：自动控制依赖，某些 API 需要它，更多缓存和异常帮助器(Source)。

用tf.function装饰函数的缺点

一般缺点

如果函数只包含少量昂贵的操作，那么加速不会有太大提升(Source)

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

没有异常捕获（应该在 Eager 模式下完成；在装饰函数之外）(Source) 调试更加困难 由于隐藏的副作用和 TF 控制流造成的限制

Detailed information on AutoGraph limitations is available.

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

不允许在tf.function 中多次创建变量，但随着 tf.v1 代码的逐步淘汰(Source) 这可能会发生变化

对于带有 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 的更改，这会导致错误，因为尚未完成跨功能分析(Source)：

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.

【讨论】：

这是一个很好的总结。还值得注意的是，当从 eager 模式调用时，tf.function 在第一次调用后有大约 200 us（给或取）的开销。不过，从另一个 tf.function 调用 tf.function 很好。所以你想包装尽可能多的计算。如果不是因为限制，你应该包装整个程序。

这个答案是 tl;dr 恕我直言，它并没有真正回答我的问题，而只是给出了我自己发现的相同的零散信息。另外，说我不应该将@tf.function 用于生产而仅用于开发并不是一个可行的解决方案。首先，在机器学习中（至少在研究中），开发阶段的训练也创造了最终产品（训练好的模型）。其次，装饰器是一个重大的变化。我不能只将它们放在“开发后”并确保代码的行为相同。这意味着我必须在那里与他们一起开发和测试。

@problemofficer 很抱歉造成混乱。在我的回答中谈论生产时，我正在考虑将培训（在大型数据集上）作为其中的一部分。在我自己的研究中，我使用一个玩具数据集在 Eager 模式下开发/调试我的函数，然后在适当的情况下添加 tf.function。

【参考方案2】：

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

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

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

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

【讨论】：

但是副作用或例如tf.range()。 AFAIK 这些无法自动转换。所以我需要从一开始就考虑使用自动图表来编写我的自定义图层。因此我不能只装饰调用（预测）函数。

【参考方案3】：

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

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

还有a section in the documentation在各种场景下讨论加速，在本节开头已经提到了上述两种情况：

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

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

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