Tensorflow：迭代数据集所需的时间比预期长 1-2 个数量级答案

【问题标题】：Tensorflow: Iterating over dataset takes 1-2 orders of magnitude longer than expectedTensorflow：迭代数据集所需的时间比预期长 1-2 个数量级
【发布时间】：2021-06-01 17:51:06
【问题描述】：

情况：

我有一个数据集 (<class 'tensorflow.python.data.ops.dataset_ops.MapDataset'> )，它是一些神经网络输出的结果。为了得到我的最终预测，我目前正在迭代它如下：

for row in dataset:
    ap_distance, an_distance = row
    y_pred.append(int(ap_distance.numpy() > an_distance.numpy()))

数据集有两列，每列包含一个用张量包裹的标量。

问题：

循环体非常简单，计算时间0.3 秒，因此从数据集中获取下一行似乎需要将近 0.3 秒！考虑到我的硬件（在具有 16 个 AMD EPYC 内核和 258GB RAM 的租用 GPU 服务器上运行）以及他笔记本电脑上的同事可以在比我少 1-2 个数量级的时间内完成迭代的事实，这种行为非常奇怪。数据集大约有。 60k 行，因此等待这么久是不可接受的。

我尝试了什么：

我尝试将上述循环体映射到数据集对象的每一行，但遗憾的是 .numpy() 在 dataset.map() 中不可用！

问题：

为什么从数据集中获取新行需要这么长时间？
如何解决此性能下降问题？

【问题讨论】：

标签： python numpy tensorflow tf.keras

【解决方案1】：

性能下降的解决方案是矢量化您的操作。

即

dataset=dataset.batch(len(dataset))
vec_ap_distance,vec_an_distance=next(iter(dataset))
y_pred=tf.cast(vec_ap_distance>vec_an_distance,tf.int32).numpy()

代码运行时间长的主要原因是有太多不必要的 numpy() 调用。

如果有 60k 行，将有 120k numpy() 调用，而在我的代码中，只能有一个 numpy() 调用。

总结：

尽可能矢量化
避免 numpy() 调用

【讨论】：