AutoML for image augmentation
Project description
AutoAlbument
AutoAlbument is an AutoML tool that learns image augmentation policies from data using the Faster AutoAugment algorithm. It relieves the user from manually selecting augmentations and tuning their parameters. AutoAlbument provides a complete ready-to-use configuration for an augmentation pipeline.
AutoAlbument supports image classification and semantic segmentation tasks.
Table of contents
Installation
AutoAlbument requires Python 3.6 or higher.
PyPI
To install the latest stable version from PyPI:
pip install -U autoalbument
GitHub
To install the latest version from GitHub:
pip install -U git+https://github.com/albumentations-team/autoalbument
Usage
1. Create a directory with configuration files.
Run autoalbument-create --config-dir </path/to/directory> --task <deep_learning_task> --num-classes <num_classes>
, e.g. autoalbument-create --config-dir ~/experiments/autoalbument-search-cifar10 --task classification --num-classes 10
.
- A value for the
--config-dir
option should contain a path to the directory. AutoAlbument will create this directory and put two files into it:dataset.py
andsearch.yaml
(more on them later). - A value for the
--task
option should contain the name of a deep learning task. Supported values areclassification
andsemantic_segmentation
. - A value for the
--num-classes
option should contain the number of distinct classes in the classification or segmentation dataset.
2. Add implementation for __len__
and __getitem__
methods in dataset.py
.
The dataset.py
file created at step 1 by autoalbument-create
contains stubs for implementing a PyTorch dataset (you can read more about creating custom PyTorch datasets here). You need to add implementation for for __len__
and __getitem__
methods (and optionally add the initialization logic if required).
A dataset for a classification task should return an image and a class label. A dataset for a segmentation task should return an image and an associated mask.
3. [Optional] Adjust search parameters in search.yaml
.
You may want to change parameters that AutoAlbument will use to search for augmentation policies. To do this, you need to edit the search.yaml
file created by autoalbument-create
at step 1. Each configuration parameter contains a comment that describes the meaning of the setting. Please refer to the "Tuning the search parameters" section that includes a description of the most critical parameters.
search.yaml
is a Hydra config file. You can use all Hydra features inside it.
4. Run a search for augmentation policies.
To search for augmentation policies, run autoalbument-search --config-dir </path/to/directory>
, e.g. autoalbument-search --config-dir ~/experiments/autoalbument-search-cifar10
. The value of --config-dir
should be the same value that was passed to autoalbument-create
at step 1.
autoalbument-search
will create a directory with output files (by default the path of the directory will be <config_dir>/outputs/<current_date>/<current_time>
, but you can customize it in search.yaml). The policy
subdirectory will contain JSON files with policies found at each search phase's epoch.
autoalbument-search
is a command wrapped with the @hydra.main
decorator from Hydra. You can use all Hydra features when calling this command.
AutoAlbument uses PyTorch to search for augmentation policies. You can speed up the search by using a CUDA-capable GPU.
5. Use found augmentation policies in your training pipeline.
AutoAlbument produces a JSON file that contains a configuration for an augmentation pipeline. You can load that JSON file with Albumentations:
import albumentations as A
transform = A.load("/path/to/policy.json")
Then you can use the created augmentation pipeline to augment the input data.
For example, to augment an image for a classification task:
transformed = transform(image=image)
transformed_image = transformed["image"]
To augment an image and a mask for a semantic segmentation task:
transformed = transform(image=image, mask=mask)
transformed_image = transformed["image"]
transformed_mask = transformed["mask"]
You can read more about using Albumentations for augmentation in those articles Image augmentation for classification, Mask augmentation for segmentation.
Refer to this section of the documentation to get examples of how to use Albumentations with PyTorch and TensorFlow 2.
Search algorithms
AutoAlbument uses the following algorithms to search for augmentation policies.
Faster AutoAugment
"Faster AutoAugment: Learning Augmentation Strategies using Backpropagation" by Ryuichiro Hataya, Jan Zdenek, Kazuki Yoshizoe, and Hideki Nakayama. Paper | Original implementation
Tuning the search parameters
search.yaml
contains parameters for the search of augmentation policies. Here is an example search.yaml
for image classification on the CIFAR-10 dataset, and here is an example search.yaml
for semantic segmentation on the Pascal VOC dataset.
Task-specific model
A task-specific model is a model that classifies images for a classification task or outputs masks for a semantic segmentation task. Settings for a task-specific model are defined by either classification_model
or semantic_segmentation_model
depending on a selected task. Ideally, you should select the same model (the same architecture and the same pretrained weights) that you will use in an actual task. AutoAlbument uses models from PyTorch Image Models and Segmentation models packages for classification and semantic segmentation respectively.
Base PyTorch parameters.
You may want to adjust the following parameters for a PyTorch pipeline:
data.dataloader
parameters such as batch_size andnum_workers
- Number of epochs to search for best augmentation policies in
optim.epochs
. - Learning rate for optimizers in
optim.main.lr
andoptim.policy.lr
.
Parameters for the augmentations search.
Those parameters are defined in policy_model
. You may want to tune the following ones:
num_sub_policies
- number of distinct augmentation sub-policies. A random sub-policy is selected in each iteration, and that sub-policy is applied to input data. The larger number of sub-policies will produce a more diverse set of augmentations. On the other side, the more sub-policies you have, the more time and data you need to tune those sub-policies correctly.num_chunks
controls the balance between speed and diversity of augmentations in a search phase. Each batch is split-up intonum_chunks
chunks, and then a random sub-policy is applied to each chunk separately. The larger the value ofnum_chunks
helps to learn augmentation policies better but simultaneously increases the searching time. Authors of FasterAutoAugment used such values fornum_chunks
that each chunk consisted of 8 to 16 images.operation_count
- the number of augmentation operations that will be applied to each input data instance. For example,operation_count: 1
means that only one operation will be applied to an input image/mask, andoperation_count: 4
means that four sequential operations will be applied to each input image/mask. The larger number of operations produces a more diverse set of augmentations but simultaneously increases the searching time.
Preprocessing transforms
If images have different sizes or you want to train a model on image patches, you could define preprocessing transforms (such as Resizing, Cropping, and Padding) in data.preprocessing
. Those transforms will always be applied to all input data. Found augmentation policies will also contain those preprocessing transforms.
Note that it is crucial for Policy Model (a model that searches for augmentation parameters) to receive images of the same size that will be used during the training of an actual model. For some augmentations, parameters depend on input data's height and width (for example, hole sizes for the Cutout augmentation).
Examples
The examples
directory contains example configs for different tasks and datasets.
FAQ
Search takes a lot of time. How can I speed it up?
Instead of a full training dataset, you can use a reduced version to search for augmentation policies. For example, the authors of Faster AutoAugment used 6000 images from the 120 selected classes to find augmentation policies for ImageNet (while the full dataset for ILSVRC contains 1.2 million images and 1000 classes).
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for autoalbument-0.0.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 113d6bb59e221efd1fee2a3c1b5c782812c27be7a81f82d7847bf3b8ed01f5f6 |
|
MD5 | 169f0cb1b22adce08ae31f6001bd2aa1 |
|
BLAKE2b-256 | 06b4101f81e1dccc81b4e2212d634e5191a6dd7004eae2f511723e753f9fe4a8 |