Improving ImageNet-1k support

W.r.t the current support for ImageNet-1k, we can improve things:

• First, let’s start leveraging TFDS. It significantly reduces the work expected to be done by a user. Let’s walk through an example.

First, the user needs to keep the ILSVRC2012_img_train.tar and ILSVRC2012_img_val.tar archives to this path: gs://[BUCKET-NAME]/tensorflow_datasets/downloads/manual.

• One this is done, the user does the following:

import tensorflow_datasets as tfds

data_dir = "gs://[BUCKET-NAME]/tensorflow_datasets"
builder = tfds.builder("imagenet2012", data_dir=data_dir)
builder.download_and_prepare()

builder.download_and_prepare() takes some time but it’s lesser than what the current process of obtaining the initial TFRecords takes.

• The the user can load the ImageNet-1k dataset with tfds.load("imagenet2012", data_dir=data_dir) and that is it.

The above two points assume the user already has access to the GCS bucket and all the necessary privileges to write data into it.

General recommendations

W.r.t

https://github.com/keras-team/keras-cv/blob/e607e05e7d73dda4fcc1f11eca11d1f71d83bee4/keras_cv/datasets/imagenet/load.py#L92

enable interleaved reading by setting num_parallel_reads=tf.data.AUTOTUNE.

W.r.t

https://github.com/keras-team/keras-cv/blob/e607e05e7d73dda4fcc1f11eca11d1f71d83bee4/keras_cv/datasets/imagenet/load.py#L113

enable prefetching of a few batches so that the accelerator doesn’t have to wait by using dataset.prefetch(tf.data.AUTOTUNE).

[sayakpaul]

Agreed with TFDS approach for simplicity.

I think it’s also possible to use local path instead of GCS bucket.

Yes, it’s possible. However, keeping things inside a GCS Bucket is necessary to leverage TPU-based training runs. So, it kind of solves different purposes.

[tanzhenyu]

tfds still requires you to download the dataset manually. Are you referring to the process of converting from .tar.gz to TFRecords?