trendingtechnology / albu_scheduler

Image data augmentation scheduler for albumentations transforms

albu_scheduler

Scheduler for albumentations transforms based on PyTorch schedulers interface

Usage

TransformMultiStepScheduler

import albumentations as A

from albu_scheduler import TransformMultiStepScheduler

transform_1 = A.Compose([
    A.RandomCrop(width=256, height=256),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
])
transform_2 = A.Compose([
    A.RandomCrop(width=128, height=128),
    A.VerticalFlip(p=0.5),
])

scheduled_transform = TransformMultiStepScheduler(transforms=[transform_1, transform_2], 
                                                  milestones=[0, 10])
dataset = Dataset(transform=scheduled_transform)

for epoch in range(100):
    train(...)
    validate(...)
    scheduled_transform.step()

TransformSchedulerOnPlateau

from albu_scheduler import TransformSchedulerOnPlateau

scheduled_transform = TransformSchedulerOnPlateau(transforms=[transform_1, transform_2], 
                                                  mode="max",
                                                  patience=5)

dataset = Dataset(transform=scheduled_transform)
for epoch in range(100):
    train(...)
    score = validate(...)
    scheduled_transform.step(score)

Installation

git clone https://github.com/KiriLev/albu_scheduler
cd albu_scheduler
make install