嵌套 tf.function 非常慢

Question

在一个用 tf.function 修饰的函数中，我尝试调用另一个用 tf.function 修饰的函数。结果慢得可怕。

那是因为我不应该在函数中使用 python 本机类型吗？

测试：

import numpy as np
import tensorflow as tf


@tf.function
def loop(x, y):
    for i in range(1000):
        x.assign_add(y)
    return x


@tf.function
def loop2(x, y):
    for i in range(1000):
        loop(x, y)
    return x


def main():
    print("TensorFlow version: {}".format(tf.__version__))
    print("Eager execution: {}".format(tf.executing_eagerly()))

    x = tf.Variable(initial_value=0, dtype=np.float32)
    y = tf.Variable(initial_value=1, dtype=np.float32)

    # print(loop2(x, y))  # horribly slow

    for i in range(1000):  # faster
        loop(x, y)


main()

Answer 1

您应该阅读您链接的答案中引用的文章 part 3。

在第 3 部分中，您可以看到问题不仅出现在使用 Python 本机类型时，而且出现在使用 Python 构造（如 for）时 Python 类型而不是 tf.Tensor 对象。

特别是，当在 range 上循环而不是在 tf.range 上循环时，您正在构建一个巨大的图表，因为您重复了 1000 次主体循环（您重新展开循环。

如果将 range 替换为 tf.range，一切都会变得更快。

证明。

您的代码（带有时间测量值和 100 而不是 1000）：

import numpy as np
import tensorflow as tf
from time import time

@tf.function
def loop(x, y):
    for i in range(100):
        x.assign_add(y)
    return x


@tf.function
def loop2(x, y):
    for i in range(100):
        loop(x, y)
    return x


def main():
    print("TensorFlow version: {}".format(tf.__version__))
    print("Eager execution: {}".format(tf.executing_eagerly()))

    x = tf.Variable(initial_value=0, dtype=np.float32)
    y = tf.Variable(initial_value=1, dtype=np.float32)
    print("one")
    start = time()
    print(loop2(x, y))  # horribly slow
    print("end: ", time() - start)
    print("second: ")
    start = time()
    for i in range(100):  # faster
        loop(x, y)
    print("end: ", time() - start)


main()

输出：

TensorFlow version: 2.0.0-beta0
Eager execution: True
one
tf.Tensor(10000.0, shape=(), dtype=float32)
end:  86.44128751754761
second: 
end:  0.08476066589355469

仅使用 TensorFlow 方法更新代码：

@tf.function
def loop__(x, y):
    for i in tf.range(100):
        x.assign_add(y)
    return x


@tf.function
def loop2__(x, y):
    for i in tf.range(100):
        loop__(x, y)
    return x


def main():
    print("TensorFlow version: {}".format(tf.__version__))
    print("Eager execution: {}".format(tf.executing_eagerly()))

    x = tf.Variable(initial_value=0, dtype=np.float32)
    y = tf.Variable(initial_value=1, dtype=np.float32)
    print("one")
    start = time()
    print(loop2__(x, y))  # horribly slow
    print("end: ", time() - start)
    print("second: ")
    start = time()
    for i in tf.range(100):  # faster
        loop__(x, y)
    print("end: ", time() - start)


main()

输出：

TensorFlow version: 2.0.0-beta0
Eager execution: True
one
tf.Tensor(10000.0, shape=(), dtype=float32)
end:  0.4946322441101074
second: 
end:  0.24096465110778809

Answer 2

关于@tf.function的几点注意事项（外行）：

@tf.function 构建它修饰的函数的可调用图
该图是通过使用作为函数签名的键来引用的。如果函数的输入是张量，则此签名为 TensorSpec；如果函数的输入不是张量，则此签名为具有参数实际值的元组
每次调用图形时，都会在所有可用的 'callable graphs' 中检查密钥，如果找到匹配项，则使用 'already built callable graph'。如果不是，则将函数转换为可调用图，然后调用它。构建图形称为通过文档跟踪功能。现在，您可以了解为什么每次调用带有 python natives 的函数都会创建一个新图形。输入的特定组合根本不作为键存在，而在张量的情况下，Tesnsorspec 键对于具有相同形状和 dtype
如果在函数内部使用 python 可迭代对象，那么 while 'tracing the function' 将展开循环以创建一个巨大的图形。如果使用像 tf.range 这样的 tensorflow 等价物，那么 tensorflow 知道如何在不展开的情况下处理这个问题。这种展开在函数第一次为运行时会产生开销，但是展开的循环总是比循环本身更快。因此，您会注意到的行为是这样的：With python iterable，与 tensorflow equivalnet (tf.range) 相反，第一个函数运行非常慢，因此图表created 将在加速器上消耗更多内存，但在所有后续运行s 上明显更快，因为具有 python 可迭代的图形使用展开循环 .

演示：

与tf.range

@tf.function
def loop__(x, y):
    for i in tf.range(10000):
        x.assign_add(y)
    return x


@tf.function
def loop2__(x, y):
    for i in tf.range(100):
        loop__(x, y)
    return x

x = tf.Variable(initial_value=0, dtype=np.float32)
y = tf.Variable(initial_value=1, dtype=np.float32)

start = time()
print(loop2__(x, y))
print("first run with tf.range", time() - start)
start = time()
print(loop2__(x, y))
print("second run with tf.range", time() - start)

output:
tf.Tensor(1000000.0, shape=(), dtype=float32)
first run with tf.range 10.322974920272827
tf.Tensor(2000000.0, shape=(), dtype=float32)
second run with tf.range 11.379822969436646

python 范围：

@tf.function
def loop__(x, y):
    for i in range(10000):
        x.assign_add(y)
    return x


@tf.function
def loop2__(x, y):
    for i in tf.range(100):
        loop__(x, y)
    return x

x = tf.Variable(initial_value=0, dtype=np.float32)
y = tf.Variable(initial_value=1, dtype=np.float32)

start = time()
print(loop2__(x, y))
print("first run with python range", time() - start)
start = time()
print(loop2__(x, y))
print("second run with python range", time() - start)

output (with loads of warnings about inefficient graph unrolling):
tf.Tensor(1000000.0, shape=(), dtype=float32)
first run with python range 51.13001751899719
tf.Tensor(2000000.0, shape=(), dtype=float32)
second run with python range 1.1093688011169434

嵌套 tf.function 非常慢

Nested tf.function is horribly slow

python

tensorflow

tensorflow2.0