How to load my own dataset

[WEIYI2021]

This repo has used cifar as the dataset for training. If I want to use my own dataset, how can I revise the code ?

[andsteing]

When you specify config.dataset, this can be a tfds dataset name (like cifar10) or the path to a directory (see the Colab vit_jax_augreg.ipynb for an illustration). The implementation of the data loading from tfds / directory is in the file vit_jax/input_pipeline.py.

If you want to use the dataset from the command line, you'll also have to extend the config, in particular this place:

vision_transformer/vit_jax/configs/common.py

Lines 110 to 115 in 60104c1

	def with_dataset(config: ml_collections.ConfigDict,
	dataset: str) -> ml_collections.ConfigDict:
	config = ml_collections.ConfigDict(config.to_dict())
	config.dataset = dataset
	config.update(DATASET_PRESETS[dataset])
	return config