albumentations-team / autoalbument

AutoML for image augmentation. AutoAlbument uses the Faster AutoAugment algorithm to find optimal augmentation policies. Documentation - https://albumentations.ai/docs/autoalbument/

Home Page:https://albumentations.ai/docs/autoalbument/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

support integrate into python code

ConstantSun opened this issue · comments

Hi, how can I run autoalbument from my code base rather than from CLI?

Hey @ConstantSun

There is a way to import and use AutoAlbument code, but this API is considered unstable for now. I am exploring options to completely rewrite the implementation using some high-level PyTorch framework such as PyTorch Lightning. The rewrite will allow us to quickly implement new features such as support for distributed training, fp16 optimizations, and TPU support.

Saying this, here is a way to use AutoAlbument in Python code. The snippet below shows how to use the CIFAR10 config:

from omegaconf import OmegaConf

from autoalbument.faster_autoaugment.search import FAAClassification, FAASemanticSegmentation

# Set the logging level so the Python code will output log messages from AutoAlbument'
import logging 
logging.basicConfig(level=logging.INFO)

# We use OmegaConf to load a YAML config file and convert it to the config object needed for AutoAlbument.
# Note that OmegaConf also supports other methods for creating config files, see more at
# https://omegaconf.readthedocs.io/en/2.0_branch/usage.html#creating
cfg = OmegaConf.load("/path/to/autoalbument/examples/cifar10/search.yaml")

# AutoAlbument requires that the dataset specified in
# https://github.com/albumentations-team/autoalbument/blob/master/examples/cifar10/search.yaml#L85
# should be located inside PYTHONPATH, so we use `sys.path.append` to add the directory that contains
# `dataset.py` into PYTHONPATH. You may need to adjust this value if you are planning to load a dataset
# from another source.
sys.path.append('/path/to/autoalbument/examples/cifar10')

# Create an instance of the search algorithm. 
# For semantic segmentation use FAASemanticSegmentation instead of FAAClassification.
searcher = FAAClassification(cfg)

# Run the actual search
searcher.search()

A few notes:

class MyFAAClassification(FAAClassification):

    def create_dataloader(self):
        dataloader = ...
        return dataloader 

searcher = MyFAAClassification(cfg)
  • By default, the code will save artifacts (such as weights and policies) to the current working directory. To change this behavior, you could change the current working directory in Python by executing os.chdir("/some/directory") before creating an instance of FAAClassification. Or, as an alternative, you could override class methods such as save_policy and save_checkpoint in your subclass.

Thank you a lot. As you said above, the code will save artifacts (such as weights and policies) to the current working directory, which model do those weights belong to? The semantic-segmentation(classification) model or the WGAN one?

All models will be saved.

Inside AutoAlbument, there are two models:

  1. Policy model (self.models['policy']). In terms of GAN, it is a Generator. This model is used to augment input images.
  2. Main model (self.models['main']). It is a classification or semantic segmentation model with an auxiliary output that predicts whether an input image is a non-augmented or augmented one. So this auxiliary output acts as a Discriminator.

Here is a method that saves a checkpoint - https://github.com/albumentations-team/autoalbument/blob/master/autoalbument/faster_autoaugment/search.py#L321
It saves both models and optimizer states for them.

Note that by default, this save_checkpoint method will be called only when the save_checkpoints config value is set to True.

Oh, I'll check them all! Thank you so much :-)))))

Hi, I tried to integrate auto-albumentations into python code like your suggestion. However, it turned out that I need to add these params in search.yaml file :
policy_model:
temperature: 0.1
operation_count: 4
task_factor : 4
My question is: What are the meanings of these above params? ( 0.1 -4 - 4 is just some number I randomly chose)

  • temperature is a parameter for Relaxed Bernoulli Distribution. During the search phase, the probabilities of applying transforms are sampled from this distribution.
  • task_factor is a coefficient for either semantic segmentation or classification loss (depending on the task). So total loss is calculated as task_factor * task_loss + WGAN_loss.
  • operation_count is a number of augmentations that will be applied to each input image. So operation_count: 1 means that only 1 augmentation will be applied, and operation_count: 4 means that 4 augmentations will be sequentially applied to an input image. Basically, the bigger the value of operation_count, the more diverse augmented images you get. However, larger values of operation_count will increase search time (because now you apply more augmentations), and those larger values require more training data to find the best augmentation policies. So the effect of increasing operation_count is similar to increasing the number of layers in a neural network: a network with more layers usually achieves better results, but it requires more data and more computational resources for training.

For temperature and task_factor I recommend to use default values (0.05 and 0.1 respectively). Please refer to the Faster AutoAugment paper if you want to tune them. In the paper, those parameters are named λ and ϵ, respectively.

For operation_count, I recommend using a value between 1 and 4, depending on your computation resources and the amount of available data.

Oh, thank you! There's one more param that I forgot to mention : gp_factor. What does this param do and what value should I assign to it?

This is a gradient penalty coefficient for training WGAN proposed in Improved Training of Wasserstein GANs. As a default, you can set gp_factor to 10.

Thanks! How can I find those default values? I couldn't find them in the autoalbuments code, do I have to read the reference papers ?

Default values for AutoAlbument are located inside Hydra code (a tool that powers AutoAlbument CLI) - https://github.com/albumentations-team/autoalbument/blob/master/autoalbument/config/faster_autoaugment.py. For now, if you use AutoAlbument inside your Python code (so without a CLI), you have to set them manually. I will explore an option to move those default values into the AutoAlbment code that powers the augmentation search.

Thank you a lot! Wish you have a good day!