python - 嵌套的 tf.function 非常慢

Question

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

那是因为我不应该在函数中使用 python 本地类型吗？ Tensorflow 2.0 model using tf.function very slow and is recompiling every time the train count changes. Eager runs about 4x faster

测试:

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 类型而非 进行操作的 Python 构造(如 for)时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