Add batchwise parameter to RandomChoice (keras-team#1749)

* Add batchwise parameter to RandomChoice * Add batchwise parameter to RandomChoice * Add batchwise parameter to RandomChoice * Add batchwise parameter to RandomChoice * blank line removed * random choice
Skinopathy-inc · May 5, 2023 · 6e1fcdd · 6e1fcdd
1 parent 0e3f121
commit 6e1fcdd
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 1 deletion.
diff --git a/keras_cv/layers/preprocessing/random_choice.py b/keras_cv/layers/preprocessing/random_choice.py
@@ -52,19 +52,26 @@ class RandomChoice(BaseImageAugmentationLayer):
             apply the augmentations. This offers a significant performance
             boost, but can only be used if all the layers provided to the
             `layers` argument support auto vectorization.
+        batchwise: Boolean, whether to pass entire batches to the
+            underlying layer. When set to `True`, each batch is passed to a
+            single layer, instead of each sample to an independent layer. This
+            is useful when using `MixUp()`, `CutMix()`, `Mosaic()`, etc.
+            Defaults to `False`.
         seed: Integer. Used to create a random seed.
     """
 
     def __init__(
         self,
         layers,
         auto_vectorize=False,
+        batchwise=False,
         seed=None,
         **kwargs,
     ):
         super().__init__(**kwargs, seed=seed, force_generator=True)
         self.layers = layers
         self.auto_vectorize = auto_vectorize
+        self.batchwise = batchwise
         self.seed = seed
 
     def _curry_call_layer(self, inputs, layer):
@@ -73,11 +80,16 @@ def call_layer():
 
         return call_layer
 
+    def _batch_augment(self, inputs):
+        if self.batchwise:
+            return self._augment(inputs)
+        else:
+            return super()._batch_augment(inputs)
+
     def _augment(self, inputs, *args, **kwargs):
         selected_op = self._random_generator.random_uniform(
             (), minval=0, maxval=len(self.layers), dtype=tf.int32
         )
-
         # Warning:
         # Do not replace the currying function with a lambda.
         # Originally we used a lambda, but due to Python's
@@ -105,6 +117,7 @@ def get_config(self):
                 "layers": self.layers,
                 "auto_vectorize": self.auto_vectorize,
                 "seed": self.seed,
+                "batchwise": self.batchwise,
             }
         )
         return config
diff --git a/keras_cv/layers/preprocessing/random_choice_test.py b/keras_cv/layers/preprocessing/random_choice_test.py
@@ -53,6 +53,26 @@ def call_pipeline(xs):
 
         self.assertAllClose(xs + 1, os)
 
+    def test_batchwise(self):
+        layer = AddOneToInputs()
+        pipeline = layers.RandomChoice(layers=[layer], batchwise=True)
+        xs = tf.random.uniform((4, 5, 5, 3), 0, 100, dtype=tf.float32)
+        os = pipeline(xs)
+
+        self.assertAllClose(xs + 1, os)
+        # Ensure the layer is only called once for the entire batch
+        self.assertEqual(layer.call_counter, 1)
+
+    def test_works_with_cutmix_mixup(self):
+        pipeline = layers.RandomChoice(
+            layers=[layers.CutMix(), layers.MixUp()], batchwise=True
+        )
+        xs = {
+            "images": tf.random.uniform((4, 5, 5, 3), 0, 100, dtype=tf.float32),
+            "labels": tf.random.uniform((4, 10), 0, 1, dtype=tf.float32),
+        }
+        pipeline(xs)
+
     def test_calls_layer_augmentation_single_image(self):
         layer = AddOneToInputs()
         pipeline = layers.RandomChoice(layers=[layer])