使用 tf.function 时，Tensorflow 逐元素梯度变慢答案

【问题标题】：Tensorflow element wise gradient slower when using tf.function使用 tf.function 时，Tensorflow 逐元素梯度变慢
【发布时间】：2021-01-22 22:45:52
【问题描述】：

我需要计算一个元素梯度作为神经网络的输入。我决定使用tf.data.Dataset 来存储输入数据。预处理数据和计算梯度是昂贵的，我想分批做，然后存储数据。

我简化了函数以使用形状 (batch_size, x, y)，并且我想分别计算每个 y 的梯度。

使用tf.GradientTape 如下所示：

import tensoflow as tf

# @tf.function
def test(inp):
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(inp)
        out = inp**2
        out = tf.unstack(out, axis=-1)
    grad = []
    for x in out:
        grad.append(tape.gradient(x, inp))
    del tape
    return tf.stack(grad, axis=-1)

inp = tf.random.normal((32, 100, 50))
test(inp)

这段代码需要 ~76 ms 来执行而不需要 3.1 s 和 tf.function 装饰器。不幸的是，当将它与 tf.data.Dataset.map 一起使用时，会出现同样的减速，我假设它会将其转换为 tf.function

我尝试改用tf.batch_jacobian，它不会受到tf.function 减速的影响，但会计算出更多的梯度，我必须减少它们。执行大约需要 15 秒。

@tf.function
def test(inp):
    with tf.GradientTape() as tape:
        tape.watch(inp)
        out = inp**2
    grad = tape.batch_jacobian(out, inp)
    return tf.math.reduce_sum(grad, axis=3)
x  = test(inp)

对于更大的数据集和更多的资源繁重的计算，我试图避免这种减速，但我还没有找到解决方案，我也不明白为什么它的计算速度要慢得多。有没有办法重塑数据并使用雅可比方法或其他东西，可以克服这个问题？

【问题讨论】：

您如何进行测量？ tf.function 必须先跟踪函数，这有一些开销，但通常是您只需支付一次的代价。

标签： python tensorflow

【解决方案1】：

让我们用 IPython 的%timeit 做一个快速实验。我定义了两个函数，一个带有 tf.function 装饰器，一个没有：

import tensorflow as tf

def test_no_tracing(inp):
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(inp)
        out = inp**2
        out = tf.unstack(out, axis=-1)
    grad = []
    for x in out:
        grad.append(tape.gradient(x, inp))
    del tape
    return tf.stack(grad, axis=-1)

@tf.function
def test_tracing(inp):
    print("Tracing")
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(inp)
        out = inp**2
        out = tf.unstack(out, axis=-1)
    grad = []
    for x in out:
        grad.append(tape.gradient(x, inp))
    del tape
    return tf.stack(grad, axis=-1)

inp = tf.random.normal((32, 100, 50))

让我们看看结果：

使用 tf.function 装饰器：

In [2]: %timeit test_tracing(inp)
Tracing
2021-01-22 15:22:15.003262: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:116] None of the MLIR optimization passes are enabled (registered 2)
2021-01-22 15:22:15.076448: I tensorflow/core/platform/profile_utils/cpu_utils.cc:112] CPU Frequency: 2599990000 Hz
10.3 ms ± 579 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

没有：

In [3]: %timeit test_no_tracing(inp)
71.7 ms ± 1.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

用tf.function 修饰的函数大约快7 倍。如果只运行一次函数，它看起来会更慢，因为修饰函数具有跟踪的开销，将代码转换为图形。跟踪完成后，代码会更快。

这可以通过只运行一次函数来验证，此时它还没有被跟踪。我们可以通过告诉%timeit 只做一次循环和一次重复来做到这一点：

In [2]: %timeit -r 1 -n 1 test_tracing(inp)
Tracing
2021-01-22 15:29:47.189850: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:116] None of the MLIR optimization passes are enabled (registered 2)
2021-01-22 15:29:47.284413: I tensorflow/core/platform/profile_utils/cpu_utils.cc:112] CPU Frequency: 2599990000 Hz
4.97 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

在这里，时间要长得多，更接近您在问题中报告的内容。但是一旦完成，跟踪函数就会快很多！让我们再做一次：

In [3]: %timeit -r 1 -n 1 test_tracing(inp)
29.1 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

您可以在指南中阅读更多关于如何使用tf.function 获得更好性能的信息：Better performance with tf.function

【讨论】：