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

Question

Situation:

I have a dataset ( <class 'tensorflow.python.data.ops.dataset_ops.MapDataset'> ) which is a result of some neural network output. In order to get my final prediction, I am currently iterating over it as follows:

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

The dataset has two columns, each holding a scalar wrapped in a tensor.

Problem:

The loop body is very simple, it takes < 1e-5 seconds to compute. Sadly, one iteration takes ca. 0.3 seconds , so fetching the next row from the dataset seems to take almost 0.3s. This behavior is very weird given my hardware (running on a rented GPU server with 16 AMD EPYC cores and 258GB RAM) and the fact that a colleague on his laptop can finish an iteration in 1-2 orders of magnitude less time than I can. The dataset has ca, 60k rows. hence it is unacceptable to wait for so long.

What I tried:

I tried mapping the above loop body onto every row of the dataset object, but sadly .numpy() is not available inside dataset.map() !

Questions:

1. Why does it take so long to get a new row from the dataset?
2. How can I fix this performance degradation?

Answer 1

The solution to the performance degradation is to vectorize your operations.

ie

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()

The major reason for the long running time in your codes is that there are too many unnecessary numpy() calls.

If there are 60k rows, there will be 120k numpy() calls, while in my code, there can be only one numpy() call.

In summary:

1. vectorize as much as possible
2. avoid numpy() calls