How do I split Tensorflow datasets?

Question

I have a tensorflow dataset based on one .tfrecord file. How do I split the dataset into test and train datasets? E.g. 70% Train and 30% test?

Edit:

My Tensorflow Version: 1.8 I've checked, there is no "split_v" function as mentioned in the possible duplicate. Also I am working with a tfrecord file.

Possible duplicate of [Split inputs into training and test sets](https://stackoverflow.com/questions/41859605/split-inputs-into-training-and-test-sets) — ted, Jul 01 '18 at 17:06
Does this answer your question? [Split a dataset created by Tensorflow dataset API in to Train and Test?](https://stackoverflow.com/questions/48213766/split-a-dataset-created-by-tensorflow-dataset-api-in-to-train-and-test) — desertnaut, Dec 03 '20 at 00:17
The question was already answered years ago, but thanks for the link — Lukas Hestermeyer, Dec 10 '20 at 12:25
Also related: https://stackoverflow.com/questions/54519309/split-tfrecords-file-into-many-tfrecords-files — xdhmoore, Jan 22 '21 at 02:11

score 43 · Accepted Answer · edited Oct 08 '19 at 13:03

43

You may use Dataset.take() and Dataset.skip():

train_size = int(0.7 * DATASET_SIZE)
val_size = int(0.15 * DATASET_SIZE)
test_size = int(0.15 * DATASET_SIZE)

full_dataset = tf.data.TFRecordDataset(FLAGS.input_file)
full_dataset = full_dataset.shuffle()
train_dataset = full_dataset.take(train_size)
test_dataset = full_dataset.skip(train_size)
val_dataset = test_dataset.skip(test_size)
test_dataset = test_dataset.take(test_size)

For more generality, I gave an example using a 70/15/15 train/val/test split but if you don't need a test or a val set, just ignore the last 2 lines.

Take:

Creates a Dataset with at most count elements from this dataset.

Skip:

Creates a Dataset that skips count elements from this dataset.

You may also want to look into Dataset.shard():

Creates a Dataset that includes only 1/num_shards of this dataset.

edited Oct 08 '19 at 13:03

Kuba Beránek

338
7
14

answered Jul 01 '18 at 20:40

ted

9,150
3
47
83

7

Isn't there a randomness issue here if the dataset is much larger than the shuffle buffer size? Since samples are shuffled only within the (relatively) small buffer, this means approximately the first 70% of samples will be the training set, next 15% will be the test set, etc. If the data is ordered somehow this would introduce bias into the training results. Probably the solution is to shard the data, then shuffle it, then split it. – buzjwa Jul 15 '18 at 13:19
1

I agree. Good comment. I suppose most use cases would then simply need to shuffle the whole dataset at once but to be truly scalable you're right – ted Jul 15 '18 at 18:29
5

Note that skip actually iterates over the dataset so it can cause big latency on large dataset – Tomasz Sętkowski Mar 14 '19 at 22:04
3

I don't recommend this as train and test sets are not disjoint: it happens that the test set contains elements of the training set – xdola Apr 15 '20 at 17:38
2

You'll want to set shuffle(reshuffle_each_iteration=False). Without that, each time a new iteration of training starts, items from the training set and validation set will get shuffled into each-other. – Marc Stogaitis Sep 10 '20 at 07:54
Does take(), skip() happen randomly? so in RNN you wouldn't split dataset using take() and skip()? – haneulkim Oct 12 '20 at 08:11

Nick Lee · Answer 2 · 2019-10-18T15:09:44.857

This question is similar to this one and this one, and I am afraid we have not had a satisfactory answer yet.

Using take() and skip() requires knowing the dataset size. What if I don't know that, or don't want to find out?
Using shard() only gives 1 / num_shards of dataset. What if I want the rest?

I try to present a better solution below, tested on TensorFlow 2 only. Assuming you already have a shuffled dataset, you can then use filter() to split it into two:

import tensorflow as tf

all = tf.data.Dataset.from_tensor_slices(list(range(1, 21))) \
        .shuffle(10, reshuffle_each_iteration=False)

test_dataset = all.enumerate() \
                    .filter(lambda x,y: x % 4 == 0) \
                    .map(lambda x,y: y)

train_dataset = all.enumerate() \
                    .filter(lambda x,y: x % 4 != 0) \
                    .map(lambda x,y: y)

for i in test_dataset:
    print(i)

print()

for i in train_dataset:
    print(i)

The parameter reshuffle_each_iteration=False is important. It makes sure the original dataset is shuffled once and no more. Otherwise, the two resulting sets may have some overlaps.

Use enumerate() to add an index.

Use filter(lambda x,y: x % 4 == 0) to take 1 sample out of 4. Likewise, x % 4 != 0 takes 3 out of 4.

Use map(lambda x,y: y) to strip the index and recover the original sample.

This example achieves a 75/25 split.

x % 5 == 0 and x % 5 != 0 gives a 80/20 split.

If you really want a 70/30 split, x % 10 < 3 and x % 10 >= 3 should do.

UPDATE:

As of TensorFlow 2.0.0, above code may result in some warnings due to AutoGraph's limitations. To eliminate those warnings, declare all lambda functions separately:

def is_test(x, y):
    return x % 4 == 0

def is_train(x, y):
    return not is_test(x, y)

recover = lambda x,y: y

test_dataset = all.enumerate() \
                    .filter(is_test) \
                    .map(recover)

train_dataset = all.enumerate() \
                    .filter(is_train) \
                    .map(recover)

This gives no warning on my machine. And making is_train() to be not is_test() is definitely a good practice.

Nice answer. Suggestion: caching (with `.cache()`) each subsample will prevent for tensorflow performing a full iteration each time (the first full iteration for each subset seems unavoidable). — Javier JC, Jan 26 '20 at 06:41
I would assume this reads the entire dataset once, but are the test_dataset and train_dataset variables then also DataSets that are iterable, or are they fully loaded in memory? Asking because I have a large file that I would prefer to not load completely in RAM. — Tominator, Apr 14 '20 at 12:06
@Tominator, I am not sure. The way my example is set up, `test_dataset` being read in full before `train_dataset` is read, `train_dataset` _has to be_ fully stored in RAM for some time, especially because I tell it to shuffle only once. But, what if the reading is controlled so that `test_dataset` is read once for every 3 time `train_dataset` is read? That way, data does not _have to be_ fully stored in RAM. Is that the actual implementation? I suspect so. TF dataset (and this kind of data-pulling API in general) is designed precisely to deal with huge dataset ........ — Nick Lee, Apr 14 '20 at 13:03
However, without inspecting source code, I cannot confirm my suspicion. — Nick Lee, Apr 14 '20 at 13:03
You should not use `all` as a variable name as it overrides Python's built-in `all(enum)` function — ted, Nov 20 '20 at 23:55
It would be good to mention that a 70/20/10% split for train/val/test datasets are possible too with modulo 7. ```test_dataset = dataset.enumerate().filter(lambda x,y: x%10==7).map(lambda x,y: y) val_dataset = dataset.enumerate().filter(lambda x,y: x%10>7).map(lambda x,y: y) train_dataset = dataset.enumerate().filter(lambda x,y: x%10<7).map(lambda x,y: y)``` — Tim, Jan 11 '21 at 05:24

How do I split Tensorflow datasets?

2 Answers2

Linked