API ReferenceOther APIsAituneTorch

aitune.torch.dataloader

View as Markdown

Dataset definition to feed data for tuning.

User should provide a dataset with iterable samples, e.g.

  • list of tensors,
  • list of list of tensors,
  • list of dicts,
  • torch.utils.data.Dataset
  • HuggingFace datasets.

Example with a list of tensors:

>>> dataset = [torch.randn(3, 224, 224) for _ in range(4)] # 4 random images >>> for batch_size, args, kwargs in samples_generator(dataset, [4]): … args[0].shape # First argument is a batch of 2 images torch.Size([4, 3, 224, 224])

Example with a list of dictionaries:

>>> dataset = [{“input”: torch.randn(3, 224, 224)} for _ in range(4)] >>> for batch_size, args, kwargs in samples_generator(dataset, [4]): … kwargs[“input”].shape torch.Size([4, 3, 224, 224])

Example with dynamic shapes use our DataLoaderFactory and MinMaxRandomDataset.

>>> random_dataset = MinMaxRandomDataset(2, [InputConfig((3, 244, 244), (3, 488, 488), kwarg_name=“image”)]) >>> dataloader_factory = DataLoaderFactory(random_dataset) >>> for batch_size, args, kwargs in samples_generator(dataloader_factory, [4]): … kwargs[“image”].shape torch.Size([4, 3, 244, 244]) torch.Size([4, 3, 488, 488])

Samples are collated into a batch with a transformers default collator with additional support for strings.

We are using torch.utils.data.Dataloader internally, we feed tuning with different batch sizes of samples.

Non-batchable inputs should be handled by the user on module level.

>>> def my_module(*args, **kwargs): … print(len(kwargs[“prompt”]), kwargs[“width”], kwargs[“height”]) >>> dataset = [{“prompt”: “Hello, world!”} for _ in range(4)] >>> for batch_size, args, kwargs in samples_generator(dataset, [4]): … my_module(*args, **kwargs, width=1024, height=758) … my_module(*args, **kwargs, width=2048, height=1536) 4 1024 758 4 2048 1536

Module Contents

Classes

NameDescription
DataLoaderFactoryFactory for the torch DataLoader.
DynamicShapeDatasetEach sample of this dataset is of a different shape.
InputConfigConfiguration for an input tensor.
MinMaxRandomDatasetDataset that contains num_samples samples with random shapes given by the input configs.
SingleBatchDatasetWrapperMakes sure that set has enough samples for batch size by repeating the elements of the dataset.
_DynamicShapeDatasetWrapperSamples of different shapes cannot be batched together.

Functions

NameDescription
_list_data_collatorDefault data collator that simply concatenates the inputs.
_make_dataloader_factoryCreate a DataLoaderFactory from dataset or if DataLoaderFactory is given return it.
_map_data_collatorTransformers default data collector copied from source.
default_data_collatorDefault data collator that simply concatenates the inputs.
ensure_enough_samplesEnsures there is enough samples in the dataset to run the model for the given number of iterations.
samples_generatorGenerate samples from the dataset with the given batch sizes and number of samples.

Data

DatasetLike

API

class aitune.torch.dataloader.DataLoaderFactory(
    dataset: aitune.torch.dataloader.DatasetLike,
    collate_fn: collections.abc.Callable | None = None,
    num_workers: int = 0
)

Factory for the torch DataLoader.

We need samples with required batch sizes thus we create a DataLoader on demand during tuning, validation or benchmarking.

If you need custom collate function you can pass it to the constructor. By default we use default_data_collator that is based on transformers default collator with additional support for strings and lists.

NOTE: A batch of strings is a list. (Tokenization might be done using huggingface datasets and mapping functionality)

More customization can be done by extending DataLoaderFactory and overriding create_dataloader method.

collate_fn
Callable | None = collate_fn or default_data_collator
aitune.torch.dataloader.DataLoaderFactory.create_dataloader(
    batch_size: int
) -> torch.utils.data.DataLoader

Create a DataLoader from the configuration.

class aitune.torch.dataloader.DynamicShapeDataset()

Bases: list, Dataset

Each sample of this dataset is of a different shape.

Example of list of tensors:

>>> dataset = DynamicShapeDataset([torch.randn(10, 10), torch.randn(20, 20)]) >>> for batch_size, args, _ in samples_generator(dataset, [4, 8]): … print(batch_size, args[0].shape) 4 torch.Size([4, 10, 10]) 4 torch.Size([4, 20, 20]) 8 torch.Size([8, 10, 10]) 8 torch.Size([8, 20, 20])

Example of list of dicts:

>>> dataset = DynamicShapeDataset([{“input”: torch.randn(10, 10)}, {“input”: torch.randn(20, 20)}]) >>> for batch_size, _, kwargs in samples_generator(dataset, [4, 8]): … print(batch_size, kwargs[“input”].shape) 4 torch.Size([4, 10, 10]) 4 torch.Size([4, 20, 20]) 8 torch.Size([8, 10, 10]) 8 torch.Size([8, 20, 20])

aitune.torch.dataloader.DynamicShapeDataset.__getitem__(
    index
)

Get an item from the dataset.

aitune.torch.dataloader.DynamicShapeDataset.__len__()

Get the length of the dataset.

class aitune.torch.dataloader.InputConfig(
    min_input: torch.Size,
    max_input: torch.Size | None = None,
    dtype: torch.dtype = torch.float32,
    kwarg_name: str | None = None,
    min_value: int | float = 0,
    max_value: int | float = 1
)
Dataclass

Configuration for an input tensor.

Example basic usage:

>>> InputConfig(min_input=(3, 24, 24), max_input=(3, 48, 48), kwarg_name=“input”) InputConfig(min_input=(3, 24, 24), max_input=(3, 48, 48), dtype=torch.float32, kwarg_name=‘input’, min_value=0, max_value=1)

Example with just min_input for both min and max:

>>> InputConfig(min_input=(3, 24, 24), kwarg_name=“input2”, min_value=0, max_value=3) InputConfig(min_input=(3, 24, 24), max_input=(3, 24, 24), dtype=torch.float32, kwarg_name=‘input2’, min_value=0, max_value=3)

dtype
dtype = torch.float32
kwarg_name
str | None = None
max_input
Size | None = None
max_value
int | float = 1
min_input
Size
min_value
int | float = 0
aitune.torch.dataloader.InputConfig.__post_init__()

Post-initialization hook.

class aitune.torch.dataloader.MinMaxRandomDataset(
    num_samples: int,
    input_configs: list[aitune.torch.dataloader.InputConfig],
    include_min_max_shapes: bool = True
)

Bases: DynamicShapeDataset

Dataset that contains num_samples samples with random shapes given by the input configs.

is_dict
= self._validate_names()
samples
= self._generate_samples()
aitune.torch.dataloader.MinMaxRandomDataset.__getitem__(
    index
)

Get a sample from the dataset.

aitune.torch.dataloader.MinMaxRandomDataset.__len__()

Return the number of samples in the dataset.

aitune.torch.dataloader.MinMaxRandomDataset._generate_sample_dict(
    gen_tensor_fn: collections.abc.Callable
)
aitune.torch.dataloader.MinMaxRandomDataset._generate_sample_list(
    gen_tensor_fn: collections.abc.Callable
)
aitune.torch.dataloader.MinMaxRandomDataset._generate_samples()
aitune.torch.dataloader.MinMaxRandomDataset._get_max_tensor(
    cfg: aitune.torch.dataloader.InputConfig
)
aitune.torch.dataloader.MinMaxRandomDataset._get_min_tensor(
    cfg: aitune.torch.dataloader.InputConfig
)
aitune.torch.dataloader.MinMaxRandomDataset._get_tensor(
    cfg: aitune.torch.dataloader.InputConfig
)
aitune.torch.dataloader.MinMaxRandomDataset._validate_names()
class aitune.torch.dataloader.SingleBatchDatasetWrapper(
    dataset: aitune.torch.dataloader.DatasetLike,
    batch_size: int
)

Bases: Dataset

Makes sure that set has enough samples for batch size by repeating the elements of the dataset.

aitune.torch.dataloader.SingleBatchDatasetWrapper.__getitem__(
    index
)

Return the item at the given index.

aitune.torch.dataloader.SingleBatchDatasetWrapper.__len__()

Return the length of the dataset which is exactly batch_size.

class aitune.torch.dataloader._DynamicShapeDatasetWrapper(
    dataset: aitune.torch.dataloader.DynamicShapeDataset,
    batch_size: int
)

Bases: Dataset

Samples of different shapes cannot be batched together.

This wrapper allows creation of batches by using same sample multiple times.

Returns given dynamic shape samples batch_size times.

NOTE: Increases dataset size by batch_size times.

>>> dataset = [{“input”: torch.randn(10, 10)}, {“input”: torch.randn(20, 20)}] >>> wrapped_dataset = _DynamicShapeDatasetWrapper(dataset, 2) >>> len(wrapped_dataset) 4 >>> wrapped_dataset[0][“input”].shape torch.Size([10, 10]) >>> wrapped_dataset[1][“input”].shape torch.Size([10, 10]) >>> wrapped_dataset[2][“input”].shape torch.Size([20, 20]) >>> wrapped_dataset[3][“input”].shape torch.Size([20, 20])

aitune.torch.dataloader._DynamicShapeDatasetWrapper.__getitem__(
    index: int
) -> typing.Any

Return the item at the given index.

Parameters:

index
int

The index of the item to return.

Returns: Any

The item at the given index.

aitune.torch.dataloader._DynamicShapeDatasetWrapper.__len__() -> int

Returns larger dataset so that we can create batch of different sample shapes.

aitune.torch.dataloader._list_data_collator(
    features: list[list]
) -> list

Default data collator that simply concatenates the inputs.

Example of list of mixed tensors and numpy arrays:

>>> features = [[torch.randn(10, 10), np.zeros((15, 15))], [torch.randn(10, 10), np.zeros((15, 15))]] >>> result = _list_data_collator(features) >>> result[0].shape, result[1].shape (torch.Size([2, 10, 10]), torch.Size([2, 15, 15]))

Example of list of strings:

>>> features = [[“Hello World”], [“Hello World”]] >>> result = _list_data_collator(features) >>> result [[‘Hello World’, ‘Hello World’]]

Parameters:

features
list[list]

List of lists of samples

Returns: list

List of concatenated tensors.

Create a DataLoaderFactory from dataset or if DataLoaderFactory is given return it.

Convenience method to deal only with DataLoaderFactory.

aitune.torch.dataloader._map_data_collator(
    features: list
) -> dict

Transformers default data collector copied from source.

NOTE: This is a copy of the transformers default data collector with additional support for strings. https://github.com/huggingface/transformers/blob/8f137b242762eb9295a431ec6eb8cd9ee673daf9/src/transformers/data/data_collator.py#L127 Transformers are Apache 2.0 licensed.

Copy has been made to add string support and avoid dependency on transformers.

aitune.torch.dataloader.default_data_collator(
    batch: list
) -> tuple | list | dict

Default data collator that simply concatenates the inputs.

Returns: tuple | list | dict

List of concatenated tensors or dict of concatenated tensors.

aitune.torch.dataloader.ensure_enough_samples(
    dataset: aitune.torch.dataloader.DatasetLike | aitune.torch.dataloader.DataLoaderFactory | torch.Tensor,
    number_of_samples: int
) -> aitune.torch.dataloader.DatasetLike | aitune.torch.dataloader.DataLoaderFactory

Ensures there is enough samples in the dataset to run the model for the given number of iterations.

NOTE: This function trims the dataset to the given number of samples.

Parameters:

dataset
DatasetLike | DataLoaderFactory | torch.Tensor

The dataset to make iterable.

number_of_samples
int

The number of samples to ensure.

Returns: DatasetLike | DataLoaderFactory

The dataset with enough samples.

aitune.torch.dataloader.samples_generator(
    dataset: aitune.torch.dataloader.DatasetLike | aitune.torch.dataloader.DataLoaderFactory | torch.Tensor,
    batch_sizes: list[int] | collections.abc.Generator[int, None, None],
    max_num_batches_per_batch_size: int | None = None
) -> collections.abc.Generator[tuple[int, list[typing.Any], dict[str, typing.Any]], None, None]

Generate samples from the dataset with the given batch sizes and number of samples.

It is convenience utility to iterate over all samples with different batch sizes in one go.

>>> dataset = [{“prompt”: “Hello, world!”} for _ in range(10)] >>> for batch_size, args, kwargs in samples_generator(dataset, [4, 8], 1): … print(batch_size, len(kwargs[“prompt”])) 4 4 8 8

>>> for batch_size, args, kwargs in samples_generator(torch.randn(10, 10), [1, 2]): … print(batch_size, args[0].shape) 1 torch.Size([1, 10, 10]) 2 torch.Size([2, 10, 10])

NOTE: for each batch size we will iterate over all samples in the dataset.

Parameters:

dataset
DatasetLike | DataLoaderFactory | torch.Tensor

The dataset to generate samples from. It can be DataLoaderFactory or any dataset/iterable and even torch.Tensor. Tensor will be treated as a single sample dataset.

batch_sizes
list[int] | Generator[int, None, None]

The batch sizes to generate samples with.

max_num_batches_per_batch_size
int | NoneDefaults to None

The maximum number of batches to use for tuning per batch size.

Returns: None

A generator of tuples of batch size, args and kwargs

aitune.torch.dataloader.DatasetLike = Sequence | torch.utils.data.Dataset