FastAugment is a handy data augmentation toolkit for vision and imaging tasks put in a single efficient TensorFlow/PyTorch extension.

To render a richly augmented batch of 128 uint8 images of 224*224 pixels with all features enabled, FastAugment takes less than 1 millisecond on Tesla T4 GPU.

Building efficient imaging data pipelines is a challenge. Common preprocessing and augmentation scripts often overload the CPU, potentially bottlenecking the CPU and RAM bandwidth while underutilizing the GPU. This repository provides a GPU-centric, user-friendly alternative for image data augmentation.

• Augmentation is performed by the GPU with a minimum memory footprint

  • Every image in the batch is sampled only once; all the augmentation transformations are applied in a single pass.
  • High throughput due to the use of texture samplers (HW acceleration units commonly used in video games)
  • Minimal GPU overhead allows the CPU to handle other tasks.

• High quality augmented data

  • Augmentation transformations are randomized across the batch, i.e., every image undergoes a different transformation.
  • Accurate modeling of 3D viewpoint changes, in-plane transformations, and perspective distortions for improved generalization

• FastAugment is easy to deploy

  • Lightweight with no dependencies
  • Can replace a significant part of a typical data augmentation pipeline with a single call or be used just as another augmentation block
  • Can be used as a mapping in tf.data.Dataset processing pipeline, or as a part of a Keras model, or in any other setup as a TensorFlow operation.

FastAugment merges a few common data augmentation techniques into a single compute block. Augmentation parameters are randomly sampled per image from a set of user-defined ranges. The list of implemented transformations follows.

• Geometric transformations:

  • Translation
  • Scale
  • In-plane rotation
  • Perspective distortions (out-of-plane rotation)
  • Horizontal and vertical flipping

• Photometric transformations:

  • Brightness correction
  • Hue correction
  • Color saturation correction
  • Gamma correction
  • Color inversion

• Other common augmentation techniques:

  • CutOut (paper)
  • mixup (paper)

• Image classification-oriented feature set (so far)
• Works on NVidia GPUs only
• Channels-last (NHWC) input and output batches are only supported so far

augment() applies a set of random geometry and color transformations to images in a batch. The applied transformation differs from one image to another one in the same batch. Transformations parameters are sampled from uniform distributions of given ranges. Their default values enable some moderate amount of augmentations.

Every image is sampled only once through a bilinear interpolator.

Transformation application order:

• horizontal and/or vertical flipping,
• perspective distortion,
• in-plane image rotation and scaling,
• translation,
• gamma correction,
• hue, saturation and brightness correction,
• mixup,
• CutOut.

Using as a tf.data.Dataset mapping:

from fast_augment import augment

...

data = dataset.batch(batch_size).map(
    lambda x, y: augment(x, y, translation=0.1, rotation=30, mixup=0.5)
)

Using as a Keras layer:

from fast_augment import Augment

...

model = tf.keras.Sequential([
    ...
    Augment(training_only=True, translation=0.1, rotation=30, mixup=0.5),
    ...
])

In PyTorch, you would need to create a FastAugment instance like this:

from fast_augment_torch import FastAugment

...

augment = FastAugment(translation=0.1, rotation=30, mixup=0.5)
batch, labels = augment(batch, labels)

Note that batch has to be in channels-last (NHWC) format, which is somewhat uncommon for PyTorch.

Arguments reference for TensorFlow functional API:

FastAugment easily compiles from source code in a few seconds as shown below. Also, the dockerfiles and the respective prebuilt images on DockerHib are available for both TensorFlow and PyTorch.

To build the TensorFlow extension you only need cmake and any standard C++ compiler. Once the code is compiled, the repository root path is to be appended to PYTHONPATH environment variable to enable Python to find the extension.

A complete copy-paste recipe for linux (tested in ubuntu- and RHEL-based distributions):

git clone https://github.com/lnstadrum/fastaugment.git
cd fastaugment/tensorflow
cmake -B build
make -C build
export PYTHONPATH=$PYTHONPATH:$(pwd)

Once compiled and appended to PYTHONPATH, FastAugment is ready for use. It is a good thing to make sure that tests are passing before going further though:

python3 test.py

You can also run the example script to get some visuals:

python3 example.py

For PyTorch, you can compile and install fast_augment_torch module as follows:

git clone https://github.com/lnstadrum/fastaugment.git
cd fastaugment/pytorch
python3 setup.py install --user

Make sure everything is fine by running unitary tests:

python3 test.py

You can also run the example script to get some visuals:

python3 example.py

Input batch (from tf_flowers dataset):

In-plane scaling (up to 10%) and rotation (up to +/-10 deg):

Perspective distortions (tilting and panning, up to +/-30 deg):

Combined (scaling up to 10%, in-plane rotation up to +/-10 deg, panning and tilting up to +/-15 deg, translation up to 10%):

Hue shift up to +/-180 deg (unrealistic):

Gamma correction in [0.25, 1.75] range:

Combined (hue shift up to +/-15 deg, gamma correction in [0.75, 1.25] range, color saturation up to +/-50%, brightness correction in up to +/-30%):

CutOut with 40% to 50% crop side size:

Mixup:

• Scaling up to 10%
• In-plane rotation up to +/-10 deg
• Panning and tilting up to +/-15 deg
• Translation up to 10%
• Hue shift up to +/-15 deg
• Gamma correction in [0.75, 1.25] range
• Color saturation up to +/-50%
• Brightness correction in up to +/-30%
• CutOut for 50% of images with 40% to 50% crop side size
• Mixup

• Test in a multi-GPU setup
• Extend to semantic segmentation: enable nearest-neighbor resampling for segmentation masks