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:
- You could adjust the behavior of the search code by subclassing FAAClassification or FAASemanticSegmentation and overriding their methods. The source code for these classes are located at https://github.com/albumentations-team/autoalbument/blob/master/autoalbument/faster_autoaugment/search.py. For example, if you want to modify the code that creates data loaders, you could make the following class:
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 ofFAAClassification
. Or, as an alternative, you could override class methods such assave_policy
andsave_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:
- Policy model (
self.models['policy']
). In terms of GAN, it is a Generator. This model is used to augment input images. - 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 astask_factor * task_loss + WGAN_loss
.operation_count
is a number of augmentations that will be applied to each input image. Sooperation_count: 1
means that only 1 augmentation will be applied, andoperation_count: 4
means that 4 augmentations will be sequentially applied to an input image. Basically, the bigger the value ofoperation_count
, the more diverse augmented images you get. However, larger values ofoperation_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 increasingoperation_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!