嵌套的 tf.function 非常慢

Question

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

那是因为我不打算在函数中使用 python 本机类型吗？ 使用 tf.function 的 Tensorflow 2.0 模型非常慢，并且每次列车计数发生变化时都会重新编译。Eager 运行速度提高了约 4 倍

测试：

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

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

在第 3 部分中，您可以看到问题不仅在于使用 Python 本机类型时，还for在于使用对 Python 类型而不是tf.Tensor对象进行操作的Python 构造（如）时。

特别是，当循环 arange而不是 a 时，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 方法更新代码：

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 version: 2.0.0-beta0
Eager execution: True
one
tf.Tensor(10000.0, shape=(), dtype=float32)
end:  0.4946322441101074
second: 
end:  0.24096465110778809