Skip to content

base_augmentation

Base Augmentation API

Classes:

Functions:

  • tf_keras_map –

    Map function for TensorFlow Keras

Classes

BaseAugmentation 

BaseAugmentation(seed: int | None = None, auto_vectorize: bool = True, data_format: str | None = None, name: str | None = None, **kwargs)

BaseAugmentation acts as a base class for various custom augmentation layers. This class provides a common interface for augmenting samples and labels. In the future, we will add support for segmentation and bounding boxes.

The only method that needs to be implemented by the subclass is

  • augment_sample: Augment a single sample during training.

Optionally, you can implement the following methods:

  • augment_label: Augment a single label during training.
  • get_random_transformations: Returns a nested structure of random transformations that should be applied to the batch. This is required to have unique transformations for each sample in the batch and maintain the same transformations for samples and labels.
  • batch_augment: Augment a batch of samples and labels during training. Needed if layer requires access to all samples (e.g. CutMix).

By default, this method will coerce the input into a batch as well as a nested structure of inputs. If auto_vectorize is set to True, the augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map. Otherwise, it will use keras.ops.map which runs sequentially.

Parameters:

  • seed (int | None, default: None ) –

    Random seed. Defaults to None.

  • auto_vectorize (bool, default: True ) –

    If True, augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map. Otherwise, it will use keras.ops.map which runs sequentially. Defaults to True.

  • data_format (str | None, default: None ) –

    Data format. Defaults to None. Will use keras.backend.image_data_format() if None.

  • name (str | None, default: None ) –

    Layer name. Defaults to None.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def __init__(
    self,
    seed: int | None = None,
    auto_vectorize: bool = True,
    data_format: str | None = None,
    name: str | None = None,
    **kwargs,
):
    """BaseAugmentation acts as a base class for various custom augmentation layers.
    This class provides a common interface for augmenting samples and labels. In the future, we will
    add support for segmentation and bounding boxes.

    The only method that needs to be implemented by the subclass is

    - augment_sample: Augment a single sample during training.

    Optionally, you can implement the following methods:

    - augment_label: Augment a single label during training.
    - get_random_transformations: Returns a nested structure of random transformations that should be applied to the batch.
        This is required to have unique transformations for each sample in the batch and maintain the same transformations for samples and labels.
    - batch_augment: Augment a batch of samples and labels during training. Needed if layer requires access to all samples (e.g. CutMix).

    By default, this method will coerce the input into a batch as well as a nested structure of inputs.
    If auto_vectorize is set to True, the augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map.
    Otherwise, it will use keras.ops.map which runs sequentially.

    Args:
        seed (int | None): Random seed. Defaults to None.
        auto_vectorize (bool): If True, augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map.
            Otherwise, it will use keras.ops.map which runs sequentially. Defaults to True.
        data_format (str | None): Data format. Defaults to None. Will use keras.backend.image_data_format() if None.
        name (str | None): Layer name. Defaults to None.

    """
    super().__init__(name=name, **kwargs)

    self.generator = keras.random.SeedGenerator(seed)
    self.data_format = data_format or keras.backend.image_data_format()

    # This is needed for compatibility with tf.data.Dataset pipeline
    self._allow_non_tensor_positional_args = True
    self.built = True
    self._convert_input_args = False

    self.training = True
    self.auto_vectorize = auto_vectorize

Functions

call 
call(inputs: NestedTensorValue, training: bool = True) -> NestedTensorValue

This method will serve as the main entry point for the layer. It will handle the input formatting and output formatting.

Parameters:

  • inputs (NestedTensorValue) –

    Inputs to be augmented.

  • training (bool, default: True ) –

    Whether the model is training or not.

Returns:

  • NestedTensorValue ( NestedTensorValue ) –

    Augmented samples or labels.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
def call(self, inputs: NestedTensorValue, training: bool = True) -> NestedTensorValue:
    """This method will serve as the main entry point for the layer. It will handle the input formatting and output formatting.

    Args:
        inputs (NestedTensorValue): Inputs to be augmented.
        training (bool): Whether the model is training or not.

    Returns:
        NestedTensorValue: Augmented samples or labels.
    """
    self.training = training
    inputs, metadata = self._format_inputs(inputs)
    return self._format_outputs(self.batch_augment(inputs), metadata)
augment_sample 
augment_sample(inputs: NestedTensorValue) -> keras.KerasTensor

Augment a single sample during training.

Note

This method should be implemented by the subclass.

Args: input(NestedTensorValue): Single sample.

Returns:

  • KerasTensor ( KerasTensor ) –

    Augmented sample.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
def augment_sample(self, inputs: NestedTensorValue) -> keras.KerasTensor:
    """Augment a single sample during training.

    !!! note

            This method should be implemented by the subclass.
    Args:
        input(NestedTensorValue): Single sample.

    Returns:
        KerasTensor: Augmented sample.
    """
    return inputs[self.SAMPLES]
augment_samples 
augment_samples(inputs: NestedTensorValue) -> keras.KerasTensor

Augment a batch of samples during training.

Parameters:

  • inputs (NestedTensorValue) –

    Batch of samples.

Returns:

  • KerasTensor ( KerasTensor ) –

    Augmented batch of samples.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
149
150
151
152
153
154
155
156
157
158
def augment_samples(self, inputs: NestedTensorValue) -> keras.KerasTensor:
    """Augment a batch of samples during training.

    Args:
        inputs (NestedTensorValue): Batch of samples.

    Returns:
        KerasTensor: Augmented batch of samples.
    """
    return self._map_fn(self.augment_sample, inputs=inputs)
augment_label 
augment_label(inputs: NestedTensorValue) -> keras.KerasTensor

Augment a single label during training.

Note

Implement this method if you need to augment labels.

Parameters:

  • inputs (NestedTensorValue) –

    Single label.

Returns:

  • KerasTensor –

    keras.KerasTensor: Augmented label.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def augment_label(self, inputs: NestedTensorValue) -> keras.KerasTensor:
    """Augment a single label during training.

    !!! note

        Implement this method if you need to augment labels.

    Args:
        inputs (NestedTensorValue): Single label.

    Returns:
        keras.KerasTensor: Augmented label.
    """
    return inputs[self.LABELS]
augment_labels 
augment_labels(inputs: NestedTensorValue) -> keras.KerasTensor

Augment a batch of labels during training.

Parameters:

  • inputs (NestedTensorValue) –

    Batch of labels.

Returns:

  • KerasTensor –

    keras.KerasTensor: Augmented batch of labels.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
175
176
177
178
179
180
181
182
183
184
def augment_labels(self, inputs: NestedTensorValue) -> keras.KerasTensor:
    """Augment a batch of labels during training.

    Args:
        inputs (NestedTensorValue): Batch of labels.

    Returns:
        keras.KerasTensor: Augmented batch of labels.
    """
    return self._map_fn(self.augment_label, inputs=inputs)
get_random_transformations 
get_random_transformations(input_shape: tuple[int, ...]) -> NestedTensorValue

Generates random transformations needed for augmenting samples and labels.

Parameters:

  • input_shape (tuple[int, ...]) –

    Shape of the input (N, ...).

Returns:

  • NestedTensorValue ( NestedTensorValue ) –

    Batch of random transformations.

Note

This method should be implemented by the subclass if the layer requires random transformations.
Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
def get_random_transformations(self, input_shape: tuple[int, ...]) -> NestedTensorValue:
    """Generates random transformations needed for augmenting samples and labels.

    Args:
        input_shape (tuple[int,...]): Shape of the input (N, ...).

    Returns:
        NestedTensorValue: Batch of random transformations.

    !!! note
            This method should be implemented by the subclass if the layer requires random transformations.
    """
    return keras.ops.arange(input_shape[0])
batch_augment 
batch_augment(inputs: NestedTensorValue) -> NestedTensorValue

Handles processing entire batch of samples and labels in a nested structure. Responsible for calling augment_samples and augment_labels.

Parameters:

  • inputs (NestedTensorValue) –

    Batch of samples and labels.

Returns:

  • NestedTensorValue ( NestedTensorValue ) –

    Augmented batch of samples and labels.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def batch_augment(self, inputs: NestedTensorValue) -> NestedTensorValue:
    """Handles processing entire batch of samples and labels in a nested structure.
    Responsible for calling augment_samples and augment_labels.

    Args:
        inputs (NestedTensorValue): Batch of samples and labels.

    Returns:
        NestedTensorValue: Augmented batch of samples and labels.
    """
    samples = inputs.get(self.SAMPLES, None)
    labels = inputs.get(self.LABELS, None)
    result = {}

    transformations = self.get_random_transformations(input_shape=keras.ops.shape(samples))

    result[self.SAMPLES] = self.augment_samples(inputs={self.SAMPLES: samples, self.TRANSFORMS: transformations})

    if labels is not None:
        result[self.LABELS] = self.augment_labels(inputs={self.LABELS: labels, self.TRANSFORMS: transformations})
    # END IF

    # preserve any additional inputs unmodified by this layer.
    for key in inputs.keys() - result.keys():
        result[key] = inputs[key]
    return result
compute_output_shape 
compute_output_shape(input_shape: tuple[int, ...], *args, **kwargs) -> tuple[int, ...]

By default assumes the shape of the input is the same as the output.

Parameters:

  • input_shape (tuple[int, ...]) –

    Input shape.

Returns:

  • tuple[int, ...] –

    tuple[int,...]: Output shape.

Note

This method should be implemented by the subclass if the output shape is different from the input shape.
Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
284
285
286
287
288
289
290
291
292
293
294
295
296
def compute_output_shape(self, input_shape: tuple[int, ...], *args, **kwargs) -> tuple[int, ...]:
    """By default assumes the shape of the input is the same as the output.

    Args:
        input_shape (tuple[int,...]): Input shape.

    Returns:
        tuple[int,...]: Output shape.

    !!! note
            This method should be implemented by the subclass if the output shape is different from the input shape.
    """
    return input_shape
get_config 
get_config()

Serialize the layer configuration.

Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
298
299
300
301
302
303
304
305
306
307
308
def get_config(self):
    """Serialize the layer configuration."""
    config = super().get_config()
    config.update(
        {
            "seed": self.seed,
            "auto_vectorize": self.auto_vectorize,
            "data_format": self.data_format,
        }
    )
    return config

BaseAugmentation1D 

BaseAugmentation1D(**kwargs)

BaseAugmentation1D acts as a base class for various custom augmentation layers. This class provides a common interface for augmenting samples and labels. In the future, we will add support for segmentation and 1D bounding boxes.

The only method that needs to be implemented by the subclass is

  • augment_sample: Augment a single sample during training.

Optionally, you can implement the following methods:

  • augment_label: Augment a single label during training.
  • get_random_transformations: Returns a nested structure of random transformations that should be applied to the batch. This is required to have unique transformations for each sample in the batch and maintain the same transformations for samples and labels.
  • batch_augment: Augment a batch of samples and labels during training. Needed if layer requires access to all samples (e.g. CutMix).

By default, this method will coerce the input into a batch as well as a nested structure of inputs. If auto_vectorize is set to True, the augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map. Otherwise, it will use keras.ops.map which runs sequentially.

Example:

class NormalizeLayer1D(BaseAugmentation1D):

    def __init__(self, **kwargs):
        ...

    def augment_sample(self, inputs):
        sample = inputs["data"]
        mu = keras.ops.mean()
        std = keras.ops.std()
        return (sample - mu) / (std + self.epsilon)

x = np.random.rand(100, 3)
lyr = NormalizeLayer(...)
y = lyr(x, training=True)
Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
def __init__(self, **kwargs):
    """BaseAugmentation1D acts as a base class for various custom augmentation layers.
    This class provides a common interface for augmenting samples and labels. In the future, we will
    add support for segmentation and 1D bounding boxes.

    The only method that needs to be implemented by the subclass is

    - augment_sample: Augment a single sample during training.

    Optionally, you can implement the following methods:

    - augment_label: Augment a single label during training.
    - get_random_transformations: Returns a nested structure of random transformations that should be applied to the batch.
        This is required to have unique transformations for each sample in the batch and maintain the same transformations for samples and labels.
    - batch_augment: Augment a batch of samples and labels during training. Needed if layer requires access to all samples (e.g. CutMix).

    By default, this method will coerce the input into a batch as well as a nested structure of inputs.
    If auto_vectorize is set to True, the augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map.
    Otherwise, it will use keras.ops.map which runs sequentially.

    Example:

    ```python

    class NormalizeLayer1D(BaseAugmentation1D):

        def __init__(self, **kwargs):
            ...

        def augment_sample(self, inputs):
            sample = inputs["data"]
            mu = keras.ops.mean()
            std = keras.ops.std()
            return (sample - mu) / (std + self.epsilon)

    x = np.random.rand(100, 3)
    lyr = NormalizeLayer(...)
    y = lyr(x, training=True)
    ```
    """
    super().__init__(**kwargs)

    if self.data_format == "channels_first":
        self.data_axis = -1
        self.ch_axis = -2
    else:
        self.data_axis = -2
        self.ch_axis = -1

Functions

BaseAugmentation2D 

BaseAugmentation2D(**kwargs)

BaseAugmentation2D acts as a base class for various custom augmentation layers. This class provides a common interface for augmenting samples and labels. In the future, we will add support for segmentation and 1D bounding boxes.

The only method that needs to be implemented by the subclass is

  • augment_sample: Augment a single sample during training.

Optionally, you can implement the following methods:

  • augment_label: Augment a single label during training.
  • get_random_transformations: Returns a nested structure of random transformations that should be applied to the batch. This is required to have unique transformations for each sample in the batch and maintain the same transformations for samples and labels.
  • batch_augment: Augment a batch of samples and labels during training. Needed if layer requires access to all samples (e.g. CutMix).

By default, this method will coerce the input into a batch as well as a nested structure of inputs. If auto_vectorize is set to True, the augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map. Otherwise, it will use keras.ops.map which runs sequentially.

Example:

class NormalizeLayer2D(BaseAugmentation2D):

    def __init__(self, name=None, **kwargs):
        ...

    def augment_sample(self, inputs):
        sample = inputs["data"]
        mu = keras.ops.mean()
        std = keras.ops.std()
        return (sample - mu) / (std + self.epsilon)

x = np.random.rand(32, 32, 3)
lyr = NormalizeLayer(...)
y = lyr(x, training=True)
Source code in neuralspot_edge/layers/preprocessing/base_augmentation.py 
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
def __init__(self, **kwargs):
    """BaseAugmentation2D acts as a base class for various custom augmentation layers.
    This class provides a common interface for augmenting samples and labels. In the future, we will
    add support for segmentation and 1D bounding boxes.

    The only method that needs to be implemented by the subclass is

    - augment_sample: Augment a single sample during training.

    Optionally, you can implement the following methods:

    - augment_label: Augment a single label during training.
    - get_random_transformations: Returns a nested structure of random transformations that should be applied to the batch.
        This is required to have unique transformations for each sample in the batch and maintain the same transformations for samples and labels.
    - batch_augment: Augment a batch of samples and labels during training. Needed if layer requires access to all samples (e.g. CutMix).

    By default, this method will coerce the input into a batch as well as a nested structure of inputs.
    If auto_vectorize is set to True, the augment_sample and augment_label methods will be vectorized using keras.ops.vectorized_map.
    Otherwise, it will use keras.ops.map which runs sequentially.

    Example:

    ```python

    class NormalizeLayer2D(BaseAugmentation2D):

        def __init__(self, name=None, **kwargs):
            ...

        def augment_sample(self, inputs):
            sample = inputs["data"]
            mu = keras.ops.mean()
            std = keras.ops.std()
            return (sample - mu) / (std + self.epsilon)

    x = np.random.rand(32, 32, 3)
    lyr = NormalizeLayer(...)
    y = lyr(x, training=True)
    ```
    """
    super().__init__(**kwargs)

    if self.data_format == "channels_first":
        self.ch_axis = -3
        self.height_axis = -2
        self.width_axis = -1
    else:
        self.ch_axis = -1
        self.height_axis = -3
        self.width_axis = -2

Functions