Update Auto/RandAugment comments, README, more.

* Add a weighted choice option for RandAugment
* Adjust magnitude noise/std naming, config

README.md

@@ -69,6 +69,7 @@ Several (less common) features that I often utilize in my projects are included.
 * Training schedules and techniques that provide competitive results (Cosine LR, Random Erasing, Label Smoothing, etc)
 * Mixup (as in https://arxiv.org/abs/1710.09412) - currently implementing/testing
 * An inference script that dumps output to CSV is provided as an example
+* AutoAugment (https://arxiv.org/abs/1805.09501) and RandAugment (https://arxiv.org/abs/1909.13719) ImageNet configurations modeled after impl for EfficientNet training (https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py)
 
 ## Results
 

timm/data/auto_augment.py

@@ -1,7 +1,7 @@
-""" Auto Augment
+""" AutoAugment and RandAugment
 Implementation adapted from:
     https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py
-Papers: https://arxiv.org/abs/1805.09501 and https://arxiv.org/abs/1906.11172
+Papers: https://arxiv.org/abs/1805.09501, https://arxiv.org/abs/1906.11172, and https://arxiv.org/abs/1909.13719
 
 Hacked together by Ross Wightman
 """
@@ -288,18 +288,18 @@ class AutoAugmentOp:
             resample=hparams['interpolation'] if 'interpolation' in hparams else _RANDOM_INTERPOLATION,
         )
 
-        # If magnitude_noise is > 0, we introduce some randomness
+        # If magnitude_std is > 0, we introduce some randomness
         # in the usually fixed policy and sample magnitude from a normal distribution
-        # with mean `magnitude` and std-dev of `magnitude_noise`.
+        # with mean `magnitude` and std-dev of `magnitude_std`.
         # NOTE This is my own hack, being tested, not in papers or reference impls.
-        self.magnitude_noise = self.hparams.get('magnitude_noise', 0)
+        self.magnitude_std = self.hparams.get('magnitude_std', 0)
 
     def __call__(self, img):
         if random.random() > self.prob:
             return img
         magnitude = self.magnitude
-        if self.magnitude_noise and self.magnitude_noise > 0:
+        if self.magnitude_std and self.magnitude_std > 0:
-            magnitude = random.gauss(magnitude, self.magnitude_noise)
+            magnitude = random.gauss(magnitude, self.magnitude_std)
         magnitude = min(_MAX_LEVEL, max(0, magnitude)) # clip to valid range
         level_args = self.level_fn(magnitude, self.hparams) if self.level_fn is not None else tuple()
         return self.aug_fn(img, *level_args, **self.kwargs)
@@ -464,16 +464,32 @@ class AutoAugment:
 
 
 def auto_augment_transform(config_str, hparams):
+    """
+    Create a AutoAugment transform
+
+    :param config_str: String defining configuration of auto augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the AutoAugment policy (one of 'v0', 'v0r', 'original', 'originalr').
+    The remaining sections, not order sepecific determine
+        'mstd' -  float std deviation of magnitude noise applied
+    Ex 'original-mstd0.5' results in AutoAugment with original policy, magnitude_std 0.5
+
+    :param hparams: Other hparams (kwargs) for the AutoAugmentation scheme
+
+    :return: A PyTorch compatible Transform
+    """
     config = config_str.split('-')
     policy_name = config[0]
     config = config[1:]
     for c in config:
         cs = re.split(r'(\d.*)', c)
-        if len(cs) >= 2:
+        if len(cs) < 2:
-            key, val = cs[:2]
+            continue
-            if key == 'noise':
+        key, val = cs[:2]
-                # noise param injected via hparams for now
+        if key == 'mstd':
-                hparams.setdefault('magnitude_noise', float(val))
+            # noise param injected via hparams for now
+            hparams.setdefault('magnitude_std', float(val))
+        else:
+            assert False, 'Unknown AutoAugment config section'
     aa_policy = auto_augment_policy(policy_name, hparams=hparams)
     return AutoAugment(aa_policy)
 
@@ -498,6 +514,36 @@ _RAND_TRANSFORMS = [
 ]
 
 
+# These experimental weights are based loosely on the relative improvements mentioned in paper.
+# They may not result in increased performance, but could likely be tuned to so.
+_RAND_CHOICE_WEIGHTS_0 = {
+    'Rotate': 0.3,
+    'ShearX': 0.2,
+    'ShearY': 0.2,
+    'TranslateXRel': 0.1,
+    'TranslateYRel': 0.1,
+    'Color': .025,
+    'Sharpness': 0.025,
+    'AutoContrast': 0.025,
+    'Solarize': .005,
+    'SolarizeAdd': .005,
+    'Contrast': .005,
+    'Brightness': .005,
+    'Equalize': .005,
+    'PosterizeTpu': 0,
+    'Invert': 0,
+}
+
+
+def _select_rand_weights(weight_idx=0, transforms=None):
+    transforms = transforms or _RAND_TRANSFORMS
+    assert weight_idx == 0  # only one set of weights currently
+    rand_weights = _RAND_CHOICE_WEIGHTS_0
+    probs = [rand_weights[k] for k in transforms]
+    probs /= np.sum(probs)
+    return probs
+
+
 def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
     hparams = hparams or _HPARAMS_DEFAULT
     transforms = transforms or _RAND_TRANSFORMS
@@ -506,33 +552,60 @@ def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
 
 
 class RandAugment:
-    def __init__(self, ops, num_layers=2):
+    def __init__(self, ops, num_layers=2, choice_weights=None):
         self.ops = ops
         self.num_layers = num_layers
+        self.choice_weights = choice_weights
 
     def __call__(self, img):
-        for _ in range(self.num_layers):
+        # no replacement when using weighted choice
-            op = random.choice(self.ops)
+        ops = np.random.choice(
+            self.ops, self.num_layers, replace=self.choice_weights is None, p=self.choice_weights)
+        for op in ops:
             img = op(img)
         return img
 
 
 def rand_augment_transform(config_str, hparams):
-    magnitude = 10
+    """
-    num_layers = 2
+    Create a RandAugment transform
+
+    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
+    sections, not order sepecific determine
+        'm' - integer magnitude of rand augment
+        'n' - integer num layers (number of transform ops selected per image)
+        'w' - integer probabiliy weight index (index of a set of weights to influence choice of op)
+        'mstd' -  float std deviation of magnitude noise applied
+    Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
+    'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2
+
+    :param hparams: Other hparams (kwargs) for the RandAugmentation scheme
+
+    :return: A PyTorch compatible Transform
+    """
+    magnitude = _MAX_LEVEL  # default to _MAX_LEVEL for magnitude (currently 10)
+    num_layers = 2  # default to 2 ops per image
+    weight_idx = None  # default to no probability weights for op choice
     config = config_str.split('-')
     assert config[0] == 'rand'
     config = config[1:]
     for c in config:
         cs = re.split(r'(\d.*)', c)
-        if len(cs) >= 2:
+        if len(cs) < 2:
-            key, val = cs[:2]
+            continue
-            if key == 'noise':
+        key, val = cs[:2]
-                # noise param injected via hparams for now
+        if key == 'mstd':
-                hparams.setdefault('magnitude_noise', float(val))
+            # noise param injected via hparams for now
-            elif key == 'm':
+            hparams.setdefault('magnitude_std', float(val))
-                magnitude = int(val)
+        elif key == 'm':
-            elif key == 'n':
+            magnitude = int(val)
-                num_layers = int(val)
+        elif key == 'n':
+            num_layers = int(val)
+        elif key == 'w':
+            weight_idx = int(val)
+        else:
+            assert False, 'Unknown RandAugment config section'
     ra_ops = rand_augment_ops(magnitude=magnitude, hparams=hparams)
-    return RandAugment(ra_ops, num_layers)
+    choice_weights = None if weight_idx is None else _select_rand_weights(weight_idx)
+    return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)

timm/data/transforms.py

@@ -190,7 +190,7 @@ def transforms_imagenet_train(
         )
         if interpolation and interpolation != 'random':
             aa_params['interpolation'] = _pil_interp(interpolation)
-        if 'rand' in auto_augment:
+        if auto_augment.startswith('rand'):
             tfl += [rand_augment_transform(auto_augment, aa_params)]
         else:
             tfl += [auto_augment_transform(auto_augment, aa_params)]