AugMix#

class composer.algorithms.AugMix(severity=3, depth=- 1, width=3, alpha=1.0, augmentation_set='all')[source]#

The AugMix data augmentation technique.

AugMix (Hendrycks et al, 2020) creates width sequences of depth image augmentations, applies each sequence with random intensity, and returns a convex combination of the width augmented images and the original image. The coefficients for mixing the augmented images are drawn from a uniform Dirichlet(alpha, alpha, ...) distribution. The coefficient for mixing the combined augmented image and the original image is drawn from a Beta(alpha, alpha) distribution, using the same alpha.

This algorithm runs on on Event.FIT_START to insert a dataset transformation. It is a no-op if this algorithm already applied itself on the State.train_dataloader.dataset.

See the Method Card for more details.

Example

from composer.algorithms import AugMix
from composer.trainer import Trainer

augmix_algorithm = AugMix(
    severity=3,
    width=3,
    depth=-1,
    alpha=1.0,
    augmentation_set="all"
)
trainer = Trainer(
    model=model,
    train_dataloader=train_dataloader,
    eval_dataloader=eval_dataloader,
    max_duration="1ep",
    algorithms=[augmix_algorithm],
    optimizers=[optimizer]
)
Parameters
  • severity (int, optional) โ€“ Severity of augmentations; ranges from 0 (no augmentation) to 10 (most severe). Default: 3.

  • depth (int, optional) โ€“ Number of augmentations per sequence. -1 enables stochastic depth sampled uniformly from [1, 3]. Default: -1.

  • width (int, optional) โ€“ Number of augmentation sequences. Default: 3.

  • alpha (float, optional) โ€“ Pseudocount for Beta and Dirichlet distributions. Must be > 0. Higher values yield mixing coefficients closer to uniform weighting. As the value approaches 0, the mixing coefficients approach using only one version of each image. Default: 1.0.

  • augmentation_set (str, optional) โ€“

    Must be one of the following options as also described in augmentation_sets:

    • "all"

      Uses all augmentations from the paper.

    • "safe"

      Like "all", but excludes transforms that are part of the ImageNet-C/CIFAR10-C test sets.

    • "original"

      Like "all", but some of the implementations are identical to the original Github repository, which contains implementation specificities for the augmentations "color", "contrast", "sharpness", and "brightness". The original implementations have an intensity sampling scheme that samples a value bounded by 0.118 at a minimum, and a maximum value of \(intensity \times 0.18 + .1\), which ranges from 0.28 (intensity = 1) to 1.9 (intensity 10). These augmentations have different effects depending on whether they are < 0 or > 0 (or < 1 or > 1). "all" uses implementations of "color", "contrast", "sharpness", and "brightness" that account for diverging effects around 0 (or 1).

    Default: "all".