Tensorflow 将数据集拆分为训练和测试导致瓶颈/缓慢

Question

I have a dataset, and when I do preprocessing on it with ds = ds.map(process_path, num_parallel_calls=AUTOTUNE).prefetch(AUTOTUNE) that line executes very fast. When I then try to access one of the dataset's elements with:

for image, label in ds.take(1):
  print(image.shape)
  image = tf.squeeze(image)
  plt.imshow(image, cmap='gray')

It takes a second or two to load; that's my first question:

Does the preprocessing only get run on the dataset when an element from the dataset is accessed, and not immediately when I call ds.map(process_path,...)?

However my main issue is that when I split the dataset ds into two, training and testing, and try to access one of the elements again, it is considerably slow... Like 20x slower. I split it into two with:

test_ds_size = int(image_count * 0.2)
train_ds = ds.skip(test_ds_size)
test_ds = ds.take(test_ds_size)

I then try to access it in the same way I do as above but replacing ds with train_ds ; my second question is:

Why is this considerably slower, just from splitting it into two?

Or am I doing something very wrong...

Any help is greatly appreciated.

Answer 1

dataset.map creates a new dataset by applying map function.
Even in the loop also when you perform dataset.take() it creates a new dataset from the number specified with very less time.
After the dataset is loaded you are performing other operations which is not relevant to tf.data performance.
You can check from the below example.

import tensorflow as tf
from time import time

dataset = tf.data.Dataset.range(1, 100)
t1 = time()
dataset = dataset.map(lambda x: x + 1)
t2 = time()
print("Time taken for map : ", t2-t1)

t3 = time()
ds = dataset.take(50)
t4 = time()
list(ds.as_numpy_iterator())
print("Time taken for take() : ",t4-t3) 

Time taken for map :  0.013489961624145508
Time taken for take() :  0.0005645751953125

Now, let's see the time taken by take() after some operation.

dataset = tf.data.Dataset.range(1, 100)
t1 = time()
dataset = dataset.map(lambda x: x + 1)
t2 = time()
print("Time taken for map : ", t2-t1)

t3 = time()
ds = dataset.take(50)
list(ds.as_numpy_iterator())
t4 = time()
print("Time taken for take() after some operation : ",t4-t3)

Time taken for map :  0.00974416732788086
Time taken for take() after some operation :  0.017722606658935547 

Coming to the split of train and test data from the existing dataset can be done in the way you have specified, but it will take time since it iterates through all the elements.

The ideal way to create tf.data.Dataset for train and test is to create it separately like shownhere . Make sure you shuffle the data before to have the proper distribution of your dataset in both train and test data.