Question

对于一些自定义代码，我需要运行一个for循环以在Tensorflow 2中动态创建变量（启用了急切的执行模式）。（在我的自定义代码中，我写入变量的值将需要渐变，因此我想在for循环中跟踪计算，以便可以从autodiff中获得渐变）。我的代码可以运行，但是速度非常慢。实际上，它比在numpy中执行相同的操作要慢几个数量级。

我已隔离问题，并提供了突出显示问题的玩具代码段。修复它可以使我修复自定义代码。

import numpy as np
import tensorflow as tf
import timeit

N = int(1e5)
data = np.random.randn(N)
def numpy_func(data):
    new_data = np.zeros_like(data)
    for i in range(len(data)):
        new_data[i] = data[i]
    return new_data

def tf_func(data):
    new_data = tf.Variable(tf.zeros_like(data))
    for i in range(len(data)):
        new_data[i].assign(data[i])
    return new_data    

%timeit numpy_func(data)
%timeit tf_func(data)

此代码段的主要内容是，在for循环中，我只需要更新变量的一部分即可。在每次迭代中，要更新的片均不同。用于更新的数据在每次迭代中都是不同的（在我的自定义代码中，这是一些简单的计算的结果，这些计算取决于变量的切片，在这里我只是使用固定的数组来隔离问题。）

我正在使用Tensorflow 2，理想情况下需要在启用急切执行的情况下运行tensorflow代码，因为部分自定义操作取决于急切执行。

我是Tensorflow的新手，对于解决此问题的任何帮助，我将不胜感激。

非常感谢，最高

Answer 1

像这样使用时，TensorFlow永远不会非常快。理想的解决方案是对计算进行矢量化处理，因此它不需要您显式循环，但这取决于您要计算的确切内容（如果需要，您可以发布有关此问题的另一个问题）。但是，使用tf.function可以获得更好的性能。我对您的函数进行了一些更改，将new_data作为输出参数，因为tf.function不允许您在第一次调用后创建变量（但是实际上，如果您删除了new_data参数，它将也可以使用，因为tf.function将在全局范围内找到该变量。）

import numpy as np
import tensorflow as tf
import timeit

# Input data
N = int(1e3)
data = np.random.randn(N)

# NumPy
def numpy_func(data, new_data):
    new_data[:] = 0
    for i in range(len(data)):
        new_data[i] = data[i]

new_data = np.zeros_like(data)
%timeit numpy_func(data, new_data)
# 143 µs ± 4.41 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# TensorFlow
def tf_func(data, new_data):
    new_data.assign(tf.zeros_like(data))
    for i in range(len(data)):
        new_data[i].assign(data[i])
new_data = tf.Variable(tf.zeros_like(data))
%timeit tf_func(data, new_data)
# 119 ms ± 3.68 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# tf.function
# This is equivalent to using it as a decorator
tf_func2 = tf.function(tf_func)
new_data = tf.Variable(tf.zeros_like(data))
tf_func2(data, new_data)  # First call is slower
%timeit tf_func2(data, new_data)
# 3.55 ms ± 40.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这是在CPU上运行的，结果在GPU上可能有很大差异。在任何情况下，如您所见，tf.function的速度仍然比NumPy慢20倍以上，但比Python函数快30倍以上。

在急切执行模式下的切片分配上的Tensorflow循环

1 个答案: