artur-galstyan / jaxonloader

A dataloader, but for JAX

Jaxonloader

A dataloader that no one asked for, but here it is anyway.

Documentation

Check out the docs here

Installation

Install this package using pip like so:

pip install jaxonloader

Philosophy

The jaxonloader package is designed to be as lightweight as possible. It's effectively the same as the PyTorch dataloader except that it uses Numpy as the backend.

Roadmap

Other types of datasets can be added but not the datasets ifself. I've decided that all I really need is the dataloader and dataset classes -- nothing more.

jaxonloader's Issues

add mnist

allow using kaggle datasets

something like get_dataset(kaggle_id) and then it downloads it using kaggle cli

Torch dataloader not slow

Hi there, nice package but I think that the reason that torch dataloader was slow due to the fact that you are transforming to tensor and normalizing for torch but not for jax.

Don't load data into GPU!

Currently, if there is a GPU available, JAX will place arrays automatically into the GPU, which we don't want in our case. We want the data to live on the CPU and be loaded to GPU only when needed (i.e. during iteration time).

Automatically set the index to 0 if iteration reached the end

As the title suggests, set the index to the initial index if done

Add some pre-processing functions

The idea is to add a list of callables to the call function of the dataloader, which ideally get called before returning the row. The functions must be pure of course and take in a single row as input and output a single row with the same shape as the original row (although there are cases we might want to lift that requirement)

I'll check it out once I get my laptop!