Augmax 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 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, Augmax tries its best to fuse transformations together, so that these data-intensive operations are be minimized.
Augmax aims to implement an API similar to that of Albumentations. An augmentation pipeline is defined as a sequence of transformations, which are then randomly applied to the input images.
import jax
import augmax
transform = augmax.Chain(
augmax.RandomCrop(256, 256),
augmax.HorizontalFlip(),
augmax.Rotate(),
)
image = ...
rng = jax.random.PRNGKey(27)
transformed_image = transform(rng, image)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).
transformed_image = jax.jit(transform)(rng, image)sub_rngs = jax.random.split(rng, images.shape[0])
transformed_images = jax.jit(jax.vmap(transform))(sub_rngs, images)
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent