https://github.com/khdlr/loox
Efficiently Composable Data Augmentation on the GPU with Jax
https://github.com/khdlr/loox
Last synced: 11 months ago
JSON representation
Efficiently Composable Data Augmentation on the GPU with Jax
- Host: GitHub
- URL: https://github.com/khdlr/loox
- Owner: khdlr
- License: apache-2.0
- Created: 2021-12-17T13:03:49.000Z (over 4 years ago)
- Default Branch: master
- Last Pushed: 2022-01-19T15:55:38.000Z (over 4 years ago)
- Last Synced: 2025-07-27T08:46:27.078Z (12 months ago)
- Language: Python
- Size: 994 KB
- Stars: 0
- Watchers: 2
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Loox
[](https://pypi.org/project/loox/) [](https://loox.readthedocs.io/en/latest/?badge=latest)
Loox is an image data augmentation framework supporting efficiently-composable transformations
with support for JAX function transformations.
Its strengths are efficient execution of complex augmentation pipelines and batched data augmentation on the GPU/TPU via the use of [`jax.jit`](jax:jax-jit) and `jax.vmap`.
In existing data augmentation frameworks,
each transformation is executed separately,
leading to many unnecessary memory reads and writes that could be avoided.
In contrast, Loox tries its best to fuse transformations together,
so that these data-intensive operations are be minimized.
## Getting Started
Loox aims to implement an API similar to that of [Albumentations](https://albumentations.ai).
An augmentation pipeline is defined as a sequence of transformations,
which are then randomly applied to the input images.
```python
import jax
import loox
transform = loox.Chain(
loox.RandomCrop(256, 256),
loox.HorizontalFlip(),
loox.Rotate(),
)
image = ...
rng = jax.random.PRNGKey(27)
transformed_image = transform(rng, image)
```
## Batch-wise Augmentation on the GPU
Leveraging the JAX infrastructure,
it is possible to greatly accelerate data augmentation by using Just-in-Time compilation (`jax.jit`),
which can execute the code on the GPU, as well as batched augmentation (`jax.vmap`).
### Augmenting a single image on the GPU
```python
transformed_image = jax.jit(transform)(rng, image)
```
### Augmenting an entire batch of images on the GPU
```python
sub_rngs = jax.random.split(rng, images.shape[0])
transformed_images = jax.jit(jax.vmap(transform))(sub_rng, images)
```