Custom Data Augmentors

In this tutorial, we will guide you through the process of building custom data augmentors and integrating them into the data flow. Specifically, we will cover how to:

1. Create and configure a custom augmentor to expand or filter datasets.

2. Implement a data augmentor that works with both individual samples and batches.

3. Apply the custom augmentor to a dataset, enabling efficient data manipulation within the pipeline.

%config InlineBackend.figure_format = 'svg'

Imports

In this section, we will import all the necessary libraries and components that we need for creating and applying a custom data augmentor.

import datasets
from typing import Iterable
from hyped import DataFlow
from hyped.typing import Int, TraceIndexList
from hyped.core import BaseDataAugmentor, BaseDataAugmentorConfig, RunContext, process_mode

Creating the Dataset and Data Flow instance

In this section, we will create a sample dataset and instantiate the DataFlow object that will manage the data transformations.

# Create a sample dataset
ds = datasets.Dataset.from_dict({"x": [0, 1, 2, 3, 4]})

# Create a data flow instance
flow = DataFlow(ds.features)

Implementing the Custom Augmentor

Here we will define a custom augmentor, which will expand the dataset by generating new samples based on a configurable parameter n. This is achieved by subclassing BaseDataAugmentor.

# Define configuration for the custom augmentor
class CustomAugmentorConfig(BaseDataAugmentorConfig):
    n: int  # Number of times to replicate the input value

# Implement the custom augmentor
class CustomAugmentor(BaseDataAugmentor[CustomAugmentorConfig]):

    def process(self, ctx: RunContext, x: Int) -> Iterable[Int]:
        # Yield the value `x` repeated `n` times
        yield from (x for _ in range(self.config.n))

CustomAugmentorConfig extends BaseDataAugmentorConfig to specify the configuration for the augmentor, which includes an integer n that defines how many times to replicate each value.

CustomAugmentor is the main augmentor class, extending BaseDataAugmentor. The process method takes in an input value x and yields it n times, creating new variations of the original dataset.

Applying the Custom Augmentor

In this section, we will apply the custom augmentor to the dataset within the data flow graph and observe the results.

# Apply the custom augmentor
out = CustomAugmentor(n=3).call(x=flow.source["x"])

# Apply the augmentation to the dataset and collect the results
out_ds = flow.apply(ds, collect={"out": out})

# Display the augmented dataset
print(out_ds.to_dict())

{'out': [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4]}

The CustomAugmentor(n=3) creates an instance of the custom augmentor, where n=3 specifies that each value in the dataset will be repeated 3 times. We call the augmentor with the data (x=flow.source["x"]), and the results are applied to the dataset using flow.apply().

Batched Processing

In this section, we will extend the custom augmentor to handle batched processing. Batched processing is useful when you want to apply augmentations over multiple samples at once, increasing efficiency, especially with large datasets.

# Define configuration for the custom augmentor
class CustomAugmentorConfigV2(BaseDataAugmentorConfig):
    n: int  # Number of times to replicate the input value

# Implement the custom augmentor
class CustomAugmentorV2(BaseDataAugmentor[CustomAugmentorConfigV2]):
    @process_mode(batched=True)
    def process(self, ctx: RunContext, x: Int) -> tuple[Int, TraceIndexList]:
        # Yield the value `x` repeated `n` times
        return x * self.config.n, list(range(len(ctx.index))) * self.config.n

We define a new configuration class CustomAugmentorConfigV2, similar to the previous one, with an added integer n to specify how many times the input should be repeated.

The process method now uses the @process_mode(batched=True) decorator, indicating that this method handles batched processing. This decorator ensures that the method is capable of handling multiple samples in one call.

The method returns two outputs:

• The first output is the augmented data, where each input value x is repeated n times.

• The second output is a list of indices, which indicates the origin of each output sample. This list is required to track where each sample came from in the input data.

# Apply the custom augmentor
out = CustomAugmentorV2(n=3).call(x=flow.source["x"])

# Apply the augmentation to the dataset and collect the results
out_ds = flow.apply(ds, collect={"out": out})

# Display the augmented dataset
print(out_ds.to_dict())

{'out': [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4]}

We apply the custom augmentor exactly the same as before. Note that the output dataset is a reordered version of the output dataset from before.