Hydrax 🐉

A zero-copy multiprocess dataloader for JAX. Built for Project RedRocket.

Installation

pip install hydrax

If you want to use any of the extra modules, hydrax.image, hydrax.tqdm, or hydrax.pandas, you can use pip install hydrax[image,tqdm,pandas] (or pip install hydrax[all]).

Wheels are currently available for Linux x86_64 with CPython 3.10, 3.11, and 3.12.

Ensure you have JAX installed and working. If you install via pip, the latest version of JAX will be installed if it is not already, but jaxlib will not be.

From Source

Clone this repository, install Python development files and a C compiler, and run:

source path/to/your/venv/bin/activate
python -m build --wheel
pip install 'dist/hydrax-<...>.whl'

Documentation

Read the online documentation for the latest version at https://redhottensors.github.io/hydrax/.

For local HTML documentation, run make html in /sphinx and browse the generated Sphinx documentation in _build/html/index.html. You will need pip install furo, which should also install Sphinx.

Usage

from hydrax import Dataloader, DataGroup, TrainingBatch, ValidationBatch

def my_loader(data, arrays, seed):
    # load data from data source into arrays, optionally augmenting using 'seed'.
    # if 'seed' is None this is a data from a validation batch
    # return any additional data for the batch

if __name__ == "main":
    my_data = ...
    array_defs = {
        "array_name": ((dim, ...), numpy_dtype, jax_dtype),
        ...
    }

    train = DataGroup(batch_size, my_data[1000:], loader_arrays=array_defs)
    valid = DataGroup(batch_size, my_data[:1000], loader_arrays=array_defs)

    dataloader = Dataloader(
        my_loader,
        train,
        validation = ("epoch", 1, valid), # run validation after every epoch
        end_at = ("epoch", 5)             # run 5 epochs in total
    )

    with dataloader: # a with block is required
        # consider using hydrax.tqdm.tbatches instead of a vanilla for loop here
        for batch in dataloader:
            if isinstance(batch, TrainingBatch):
                run_training_batch(batch)
            elif isinstance(batch, ValidationBatch):
                run_validation_batch(batch)

            del batch # important, release batch before waiting for next one or cleaning up

# with hydrax.tqdm.tbatches
    from hydrax.tdqm import tbatches

    for batch in tbatches(dataloader, report_interval=1000): # tbatches includes a with block for the dataloader
        ...
        del batch # important, see above

Deadlocks / Stalls

If you are experiencing deadlocks as a result of retaining batch or array references between iterations, consider using debug_batch_references or gc.get_referrers to find out what's holding on to your batches, though do keep in mind that JAX dispatch will retain references while running ahead. You can check your work by running the Dataloader with depth = 1, which will immediately deadlock if the first batch is not properly released.

Batch Structure

In Hydrax, a single Dataloader is usually responsible for producing both your training and validation batches, in order to conserve resources and ensure perfectly smooth loading throughout.

Each batch produced by the Dataloader is either a TrainingBatch instance or a ValidationBatch instance, which both inherit the common functionality of Batch. (You can click any of the preceding links to view the online documentation.)

The most important properties of a Batch are:

• arrays -- { 'array_name': jax.Array, ... }, corresponding to each array defined by the source DataGroup. The first dimension of the array is the batch size.
• additional -- { 'key': [item_0_value, ...] }, corresponding to any additional data returned by your loader function. Each list's len is the batch size. If no corresponding item was returned, the element is None. Use get_additional(key[, index]) if your loader sometimes omits returning certain keys.
• data -- A proxy type to the original data descriptors for each item, with length equal to the batch size.

As mentioned above, remember to release any references to a batch or its arrays as soon as you're done with them.

Loader Processes

Read the documentation for loader_func carefully. If you receive a warning from Hydrax about your loader, you should fix your code. Failure to do this could result in your batch data changing out from underneath you, leading to significant training issues such as NaNs.

Do not attempt to construct a Dataloader inside a loader process. Ensure your training code is guarded with if __name__ == '__main__':, or is otherwise prevented from running. As a last resort, you can check hydrax.is_worker and bail.

KeyboardInterrupt (Ctrl+C / SIGINT)

The Dataloader installs a handler for KeyboardInterrupt (Ctrl+C / SIGINT) which stops the flow of batches as soon as possible. After the dataloader has completed, you can check if this occurred by reading its interrupted property. You may want to save a checkpoint along with the numbers of the current epoch and batch, so that you can resume from where you left off with start_at.

If you send a second KeyboardInterrupt, Hydrax will raise a KeyboardInterrupt at the beginning of the next batch. This exception may cause you to lose progress unless you or a framework takes care to save a checkpoint in response.

If you send a third KeyboardInterrupt, the Python interpreter is immediately stopped and control is returned to you. You will lose all progress since the last checkpoint.

Compatibility

A convienient wrapper with tqdm progress bars is provided in hydrax.tqdm. The corresponding extra is tqdm.

ICC-profile-aware 8bbp image loading with Pillow is provided in hydrax.image, and support is included for Oklab as well. The corresponding extra is image.

Compatibility for Pandas datasets is provided by hydrax.pandas. The corresponding extra is pandas.

License

Hydrax is available under the terms of the Mozilla Public License, version 2.0.