Different ways of splitting tensorflow dataset

tf_dataset_split_util.py

def split_dataset(ds, version=1):
    if version == 1:
        train_ds = ds.dataset.shard(num_shards=4, index=0)
        train_ds.concatenate(ds.dataset.shard(num_shards=4, index=1))
        train_ds.concatenate(ds.dataset.shard(num_shards=4, index=2))
        valid_ds = ds.dataset.shard(num_shards=4, index=3)
        return train_ds, valid_ds

    elif version == 2:
        def is_val(x, y):
            return x % 20 == 0

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

        recover = lambda x, y: y

        train_ds = ds.dataset.enumerate().filter(is_train).map(recover)
        valid_ds = ds.dataset.enumerate().filter(is_val).map(recover)
        return train_ds, valid_ds

    elif version == 3:
        train_ds = ds.dataset.take(ds.n_samples - 1000)
        valid_ds = ds.dataset.skip(ds.n_samples - 1000)
        return train_ds, valid_ds