时间戳

首页 > 解决方案 > 使用 tf.function 时，Tensorflow 逐元素梯度变慢

python - 使用 tf.function 时，Tensorflow 逐元素梯度变慢

问题描述

我需要计算一个元素梯度作为神经网络的输入。我决定使用 atf.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)

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

标签： pythontensorflow

解决方案

因此，让我们用 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您可以在指南中阅读有关如何获得更好性能的更多信息：使用 tf.function获得更好的性能

推荐阅读