在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热切的代码转换为与图形兼容的TensorFlow操作。这包括控制流，如if、for、while。

tf.function对TensorFlow操作最有效，NumPy和NumPy调用被转换为常量。

请参阅下面所示的代码，以便在@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的信息，请参考这。