Is it possible to only load part of a TensorFlow dataset?

Viewed 1954

I have a notebook in Google Colab with the following code:

batch_size = 64
dataset_name = 'coco/2017_panoptic'

tfds_dataset, tfds_info = tfds.load(
    dataset_name, 
    split='train', 
    with_info=True)

I would like to know if it possible to only download part of the dataset (say: 5%, or X number of images) with the tfds_load function. As far as I can see in the documentation, there are no arguments to do so. Of course it would be possible to slice the dataset after dowloading, but this particular dataset (coco/2017_panoptic) is 19.57 GiB, which obviously takes quite a while to download.

2 Answers

The original question was about how to download a subset of the dataset.

And so the answer recommending the use of an argument like split='train[:5%]' as a way of downloading only 5% of the training data is mistaken. It seems that this still downloads the entire dataset, but then only loads 5%.

You can check this for yourself by running mnist_ds_5p = tfds.load("mnist", split="train[:5%]") followed by mnist_ds = tfds.load("mnist", split="train")

No downloading takes place after running the second command. This is because the entire dataset has already been downloaded and cached after running the first command!

As many of the datasets are being fetched from a compressed form I doubt there is a simple way to avoid downloading the entire dataset I'm afraid.

Yes, it is; from the documentation:

split Which split of the data to load (e.g. 'train', 'test', ['train', 'test'], 'train[80%:]',...). See our split API guide. If None, will return all splits in a Dict[Split, tf.data.Dataset]

So, if I want the first 5% of the training set of the MNIST dataset, I could do it that way :

>>> mnist_ds_5p = tfds.load("mnist", split="train[:5%]")
>>> # looking at the size of the dataset
>>> mnist_ds_5p.cardinality()
<tf.Tensor: shape=(), dtype=int64, numpy=3000>

The MNIST dataset contains 60000 training images, and we have 3000, so 5%.

You can read more about the splitting syntax by reading the Splits and slicing guide

Related