helios.data.datamodule

Attributes

DATASET_REGISTRY	Global instance of the registry for datasets.
COLLATE_FN_REGISTRY	Global instance of the registry for collate functions.

Classes

DatasetSplit	The different dataset splits.
DataLoaderParams	Params used to create the dataloader object.
Dataset	The dataset and corresponding data loader params.
DataModule	Base class that groups together the creation of the main training datasets.

Functions

create_dataset(→ torch.utils.data.Dataset)	Create a dataset of the given type.
create_collate_fn(→ Callable)	Create a collate function of the given type.
create_dataloader(, batch_size, shuffle, num_workers, ...)	Create the dataloader for the given dataset.

Module Contents

helios.data.datamodule.DATASET_REGISTRY

Global instance of the registry for datasets.

Example

import helios.data as hld

# This automatically registers your dataset.
@hld.DATASET_REGISTRY.register
class MyDataset:
    ...

# Alternatively you can manually register a dataset like this:
hld.DATASET_REGISTRY.register(MyDataset)

helios.data.datamodule.COLLATE_FN_REGISTRY

Global instance of the registry for collate functions.

Example

import helios.data as hld

# This automatically registers your collate function.
@hld.COLLATE_FN_REGISTRY
def my_collate_fn():
    ...

# Alternatively you can manually register a collate function like this:
hld.COLLATE_FN_REGISTRY.register(my_collate_fn)

helios.data.datamodule.create_dataset(type_name: str, *args: Any, **kwargs: Any) → torch.utils.data.Dataset

Create a dataset of the given type.

This uses DATASET_REGISTRY to look-up dataset types, so ensure your datasets have been registered before using this function.

Parameters:

• type_name – the type of the dataset to create.

• args – positional arguments to pass into the dataset.

• kwargs – keyword arguments to pass into the dataset.

Returns:

The constructed dataset.

helios.data.datamodule.create_collate_fn(type_name: str, *args: Any, **kwargs: Any) → Callable

Create a collate function of the given type.

This uses COLLATE_FN_REGISTRY to look-up collate functions, so ensure that your functions have been registered before using this function. In order to support regular functions as well as callable objects, this function behaves as follows:

1. If type_name refers to a function-type, then args and kwargs are

   ignored.

2. If type_name refers to an actual type (i.e. a class), then args and

   kwargs are forwarded upon instantiation.

Parameters:

• type_name – the type of the function to create.

• args – positional arguments to pass into the collate class

• kwargs – keyword arguments to pass into the collate class

Returns:

The function

class helios.data.datamodule.DatasetSplit

Bases: enum.Enum

The different dataset splits.

TRAIN = 0

VALID = 1

TEST = 2

static from_str(label: str) → DatasetSplit

Convert the given string to the corresponding enum value.

Must be one of “train”, “test”, or “valid”

Parameters:

label – the label to convert.

Returns:

The corresponding enum value.

Raises:

ValueError – if the given value is not one of “train”, “test”, or “valid”.

helios.data.datamodule.create_dataloader(dataset: torch.utils.data.Dataset, random_seed: int = rng.get_default_seed(), batch_size: int = 1, shuffle: bool = False, num_workers: int = 0, pin_memory: bool = False, drop_last: bool = False, is_distributed: bool = False, sampler: helios.data.samplers.ResumableSamplerType | None = None, collate_fn: Callable | None = None, prefetch_factor: int | None = None, persistent_workers: bool = False, pin_memory_device: str = '', timeout: float = 0, multiprocessing_context: Any | None = None) → tuple[torch.utils.data.DataLoader, helios.data.samplers.ResumableSamplerType]

Create the dataloader for the given dataset.

If no sampler is provided, the choice of sampler will be determined based on the values of is_distributed and shuffle. Specifically, the following logic is used:

• If is_distributed, then sampler is ResumableDistributedSampler.

• Otherwise, if shuffle then sampler is ResumableRandomSampler, else ResumableSequentialSampler.

You may override this behaviour by providing your own sampler instance.

Warning

If you provide a custom sampler, then it must be derived from one of helios.data.samplers.ResumableSampler or helios.data.samplers.ResumableDistributedSampler.

Parameters:

• dataset – the dataset to use.

• random_seed – (optional) seed for the worker processes. Defaults to the value returned by get_default_seed().

• batch_size – (optional) number of samples per batch.

• shuffle – (optional) if true, samples are randomly shuffled.

• num_workers – (optional) number of worker processes for loading data.

• pin_memory – (optional) if true, use page-locked device memory.

• drop_last – (optional) if true, remove the final batch.

• is_distributed – (optional) if true, create the distributed sampler.

• sampler – (optional) sampler to use.

• collate_fn – (optional) function to merge batches.

• prefetch_factor – (optional) number of batches to prefetch per worker. Only valid when num_workers > 0.

• persistent_workers – (optional) if true, keep worker processes alive between epochs. Only valid when num_workers > 0.

• pin_memory_device – (optional) target device for pinned memory when pin_memory is true.

• timeout – (optional) timeout in seconds for collecting a batch from workers.

• multiprocessing_context – (optional) method for spawning worker processes (e.g. "fork", "spawn", "forkserver").

Returns:

The dataloader and sampler.

Raises:

TypeError – if sampler is not None and not derived from one of ResumableDistributedSampler or ResumableSampler.

class helios.data.datamodule.DataLoaderParams

Params used to create the dataloader object.

Parameters:

• random_seed – (optional) seed for the worker processes. Defaults to the value returned by get_default_seed().

• batch_size – (optional) number of samples per batch.

• shuffle – (optional) if true, samples are randomly shuffled.

• num_workers – (optional) number of worker processes for loading data.

• pin_memory – (optional) if true, use page-locked device memory.

• drop_last – (optional) if true, remove the final batch.

• is_distributed – (optional) if true, create the distributed sampler.

• sampler – (optional) sampler to use.

• collate_fn – (optional) function to merge batches.

• prefetch_factor – (optional) number of batches to prefetch per worker. Only valid when num_workers > 0.

• persistent_workers – (optional) if true, keep worker processes alive between epochs. Only valid when num_workers > 0.

• pin_memory_device – (optional) target device for pinned memory when pin_memory is true.

• timeout – (optional) timeout in seconds for collecting a batch from workers.

• multiprocessing_context – (optional) method for spawning worker processes (e.g. "fork", "spawn", "forkserver").

random_seed: int = 6691

batch_size: int = 1

shuffle: bool = False

num_workers: int = 0

pin_memory: bool = False

drop_last: bool = False

is_distributed: bool | None = None

sampler: helios.data.samplers.ResumableSamplerType | None = None

collate_fn: Callable | None = None

prefetch_factor: int | None = None

persistent_workers: bool = False

pin_memory_device: str = ''

timeout: float = 0

multiprocessing_context: Any | None = None

to_dict() → dict[str, Any]

Convert the params object to a dictionary using shallow copies.

classmethod from_dict(table: dict[str, Any])

Create a new params object from the given table.

class helios.data.datamodule.Dataset

The dataset and corresponding data loader params.

Parameters:

• dataset – the dataset.

• params – the data loader params.

dataset: torch.utils.data.Dataset

params: DataLoaderParams

dict() → dict[str, Any]

Convert to a dictionary.

class helios.data.datamodule.DataModule

Bases: abc.ABC

Base class that groups together the creation of the main training datasets.

The use of this class is to standardize the way datasets and their respective dataloaders are created, thereby allowing consistent settings across models.

Example

from torchvision.datasets import CIFAR10
from helios import data
from helios.data import Dataset, DataLoaderParams

class MyDataModule(data.DataModule):
    def prepare_data(self) -> None:
        # Use this function to prepare the data for your datasets. This will
        # be called before the distributed processes are created (if using)
        # so you should not set any state here.
        CIFAR10(download=True) # download the dataset only.

    def setup(self) -> None:
        # Register the training phase(s) using _add_train_phase. The first
        # call also sets self._train_dataset automatically.
        self._add_train_phase(CIFAR10(train=True), DataLoaderParams(...))

        # For multi-phase training, call _add_train_phase again:
        self._add_train_phase(
            CIFAR10(train=True, transform=...), DataLoaderParams(...)
        )

        # Validation and testing datasets use their own helpers.
        # A dict of settings is also accepted in place of DataLoaderParams.
        settings = {"batch_size": 1, ...}
        self._add_valid_dataset(CIFAR10(train=False), settings)
        self._add_test_dataset(CIFAR10(train=False), DataLoaderParams(...))

    def teardown(self) -> None:
        # Use this function to clean up any state. It will be called after
        # training is done.

property is_distributed: bool

Flag controlling whether distributed training is being used or not.

property trainer: helios.trainer.Trainer

Reference to the trainer.

property train_dataset: torch.utils.data.Dataset | None

The training dataset (if available).

property valid_dataset: torch.utils.data.Dataset | None

The validation dataset (if available).

property test_dataset: torch.utils.data.Dataset | None

The testing dataset (if available).

prepare_data() → None

Prepare data for training.

This can include downloading datasets, preparing caches, or streaming them from external services. This function will be called on the primary process when using distributed training (will be called prior to initialization of the processes) so don’t store any state here.

abstractmethod setup() → None

Construct all required datasets.

train_dataloader() → tuple[torch.utils.data.DataLoader, helios.data.samplers.ResumableSamplerType] | None

Create the train dataloader (if available).

valid_dataloader() → tuple[torch.utils.data.DataLoader, helios.data.samplers.ResumableSamplerType] | None

Create the valid dataloader (if available).

test_dataloader() → tuple[torch.utils.data.DataLoader, helios.data.samplers.ResumableSamplerType] | None

Create the test dataloader (if available).

teardown() → None

Clean up any state after training is over.

get_train_steps_per_epoch() → int

Return the number of iterations per training epoch.

The number is determined by constructing the training dataloader and returning its length. This function is generally useful for initialising schedulers that need the total number of steps per epoch, so you should call this within steup() or setup() and cache the result.

This function is equivalent to:

import helios.core as hlc
# If called from within Model.setup()
dataloader, _ = hlc.get_from_optional(
        self.trainer.datamodule.train_dataloader())
steps = len(dataloader)

Note

This function only applies to training datasets. If you need something similar for validation or testing, you can obtain the dataloader directly.

Returns:

The number of iterations per training epoch.

Raises:

RuntimeError – if the training dataloader hasn’t been created.

advance_train_phase() → None

Advance the training dataset to the next phase.

If the current phase is already the last one, this function does nothing.

state_dict() → dict[str, Any]

Get the full state dictionary of the datamodule.

The contents of the dictionary depend on whether training phases have been registered: * If they have, then this returns the current phase. * If they haven’t, then this returns an empty dictionary.

Returns:

The state dictionary of the datamodule.

load_state_dict(state_dict: dict[str, Any]) → None

Load the datamodule state from the given state dictionary.

Parameters:

state_dict – the state dictionary to load from.