循环 tf.data.Dataset 非常慢

Question

我想知道为什么对 tf.data.Dataset 样本的 for 循环比在相应的 numpy 数组上循环慢得多。

import numpy as np
import tensorflow as tf
import time

a = np.ones(100000, dtype=np.float32)

start_time = time.time()
for x in a:
    pass
print(time.time() - start_time)

start_time = time.time()
for x in tf.data.Dataset.from_tensor_slices(a):
    pass
print(time.time() - start_time)

---

0.05548405647277832
5.67711615562439

我的 TensorFlow 版本是 2.0.0。

Answer 1

是的，即使我也观察到相同的行为。要提高速度/性能，请尝试将 tf.data.dataset 包装在 @tf.function 中，这将花费几乎相同的时间。

AutoGraph 在tf.function 中是默认设置，并将您的 Python Eager 代码转换为与图形兼容的 TensorFlow 操作。这包括控制流，如if、for、while。

tf.function 最适用于 TensorFlow ops，NumPy 和 Python 调用被转换为常量。

请参考下面显示的代码以包裹在@tf.function中

@tf.function

def oper(a):
    start_time = time.time()
    for x in tf.data.Dataset.from_tensor_slices(a):
        pass
    print(time.time() - start_time)

下面显示numpy 和tf.data.dataset 性能之间的完整工作代码

import numpy as np
import tensorflow as tf
import time

a = np.ones(100000, dtype=np.float32)

start_time = time.time()
for x in a:
    pass
print(time.time() - start_time)


@tf.function

def oper(a):
    start_time = time.time()

    for x in tf.data.Dataset.from_tensor_slices(a):
        pass
    print(time.time() - start_time)

oper(a) 

输出：

0.012496232986450195
0.017792224884033203

要了解更多关于tf.function的信息，请参考this。