Efficient vision data augmentations for CPU/GPU per-sample/batched data.

This repository is now archived as I work on other projects. I do not plan to reopen it any time soon.

Torchaug

Introduction

Torchaug is a data augmentation library for the Pytorch ecosystem. It is meant to deal efficiently with tensors that are either on CPU or GPU and either per sample or on batches.

It enriches Torchvision (v2) that has been implemented over Pytorch and Pillow to, among other things, perform data augmentations. Because it has been implemented first with per-sample CPU data augmentations in mind, it has several drawbacks to make it efficient:

• For data augmentations on GPU, some CPU/GPU synchronizations cannot be avoided.
• For data augmentations applied on batch, the randomness is sampled for the whole batch and not each sample.

Torchaug removes these issues and its transforms are meant to be used in place of Torchvision. It is based on the code base of Torchvision and therefore follows the same nomenclature as Torchvision with functional augmentations and transforms class wrappers. However, Torchaug does not support transforms on Pillow images.

More details can be found in the documentation.

To be sure to retrieve the same data augmentations as Torchvision, the components are tested to match Torchvision outputs. We made a speed comparison here.

If you find any unexpected behavior or want to suggest a change please open an issue.

How to use

1. Install Torchaug.

pip install torchaug

2. Import data augmentations from the torchaug.transforms package just as for Torchvision.

from torchaug.transforms import (
    RandomColorJitter,
    RandomGaussianBlur,
    SequentialTransform
)


transform = SequentialTransform([
    RandomColorJitter(...),
    RandomGaussianBlur(...)
])

For a complete list of transforms please see the documentation.

How to contribute

Feel free to contribute to this library by making issues and/or pull requests. For each feature you implement, add tests to make sure it works. Also, please update the documentation.

Credits

We would like to thank the authors of Torchvision for generously opening their source code. Portions of Torchaug were originally taken from Torchvision, which is released under the BSD 3-Clause License. Please see their repository and their BSD 3-Clause License for more details.

LICENSE

Torchaug is licensed under the CeCILL-C license.