Package: jax-dataloader
Dataloader for JAXOverviewjax_dataloaderbringspytorch-likedataloader API tojax. It
supports4 datasets to download and pre-process data:jax datasethuggingface datasetspytorch
Datasettensorflow dataset3 backends to iteratively load batches:jax
dataloaderpytorch
dataloadertensorflow datasetA minimumjax-dataloaderexample:importjax_dataloaderasjdldataloader=jdl.DataLoader(dataset,# Can be a jdl.Dataset or pytorch or huggingface datasetbackend='jax',# Use 'jax' for loading data (also supports `pytorch`))batch=next(iter(dataloader))# iterate next batchInstallationThe latestjax-dataloaderrelease can directly be installed from PyPI:pipinstalljax-dataloaderor install directly from the repository:pipinstallgit+https://github.com/BirkhoffG/jax-dataloader.gitNoteWe keepjax-dataloader’s dependencies minimum, which only installjax-related dependencies, andplum-dispatchfor backend
dispatching. If you wish to use integration ofpytorch, huggingfacedatasets, ortensorflow, we recommend manually install those
dependencies.You can also runpip install jax-dataloader[all]to install
everything (not recommended).Usagejax_dataloader.core.DataLoaderfollows similar API as the pytorch dataloader.Thedatasetshould be an object of the subclass ofjax_dataloader.core.Datasetortorch.utils.data.Datasetor (the
huggingface)datasets.Datasetortf.data.Dataset.Thebackendshould be one of"jax"or"pytorch"or"tensorflow". This argument specifies which backend dataloader to
load batches.Note that not every dataset is compatible with every backend. See the
compatibility table below:jdl.Datasettorch_data.Datasettf.data.Datasetdatasets.Dataset"jax"✅❌❌✅"pytorch"✅✅❌✅"tensorflow"✅❌✅✅UsingArrayDatasetThejax_dataloader.core.ArrayDatasetis an easy way to wrap multiplejax.numpy.arrayinto one Dataset. For example, we can create anArrayDatasetas follows:# Create features `X` and labels `y`X=jnp.arange(100).reshape(10,10)y=jnp.arange(10)# Create an `ArrayDataset`arr_ds=jdl.ArrayDataset(X,y)Thisarr_dscan be loaded byeverybackends.# Create a `DataLoader` from the `ArrayDataset` via jax backenddataloader=jdl.DataLoader(arr_ds,'jax',batch_size=5,shuffle=True)# Or we can use the pytorch backenddataloader=jdl.DataLoader(arr_ds,'pytorch',batch_size=5,shuffle=True)Using Huggingface DatasetsThe huggingfacedatasetsis a
morden library for downloading, pre-processing, and sharing datasets.jax_dataloadersupports directly passing the huggingface datasets.fromdatasetsimportload_datasetFor example, We load the"squad"dataset fromdatasets:hf_ds=load_dataset("squad")Then, we can usejax_dataloaderto load batches ofhf_ds.# Create a `DataLoader` from the `datasets.Dataset` via jax backenddataloader=jdl.DataLoader(hf_ds['train'],'jax',batch_size=5,shuffle=True)# Or we can use the pytorch backenddataloader=jdl.DataLoader(hf_ds['train'],'pytorch',batch_size=5,shuffle=True)Using Pytorch DatasetsThepytorch Datasetand its
ecosystems (e.g.,torchvision,torchtext,torchaudio) supports many
built-in datasets.jax_dataloadersupports directly passing the
pytorch Dataset.NoteUnfortuantely, thepytorch
Datasetcan only work withbackend=pytorch. See the belowing example.fromtorchvision.datasetsimportMNISTimportnumpyasnpWe load the MNIST dataset fromtorchvision. TheToNumpyobject
transforms images tonumpy.array.classToNumpy(object):def__call__(self,pic):returnnp.array(pic,dtype=float)pt_ds=MNIST('/tmp/mnist/',download=True,transform=ToNumpy(),train=False)Thispt_dscanonlybe loaded via"pytorch"dataloaders.dataloader=jdl.DataLoader(pt_ds,'pytorch',batch_size=5,shuffle=True)Using Tensowflow Datasetsjax_dataloadersupports directly passing thetensorflow
datasets.importtensorflow_datasetsastfdsimporttensorflowastfFor instance, we can load the MNIST dataset fromtensorflow_datasetstf_ds=tfds.load('mnist',split='test',as_supervised=True)and usejax_dataloaderfor iterating the dataset.dataloader=jdl.DataLoader(tf_ds,'tensorflow',batch_size=5,shuffle=True)