循环 tf.data.Dataset 非常慢

我想知道为什么对 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。

是的，即使我也观察到了同样的行为。要提高速度/性能，请尝试包装tf.data.dataset in a @tf.函数几乎需要同样的时间。

签名默认为tf.function并将您的 Python eager 代码转换为与图兼容的 TensorFlow 操作。这包括控制流，例如if, for, while.

tf.function最适合 TensorFlow 操作，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 and 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) 

output:

0.012496232986450195
0.017792224884033203

想了解更多tf.function请参考this.