helios.data.datamodule ====================== .. py:module:: helios.data.datamodule Attributes ---------- .. autoapisummary:: helios.data.datamodule.DATASET_REGISTRY helios.data.datamodule.COLLATE_FN_REGISTRY Classes ------- .. autoapisummary:: helios.data.datamodule.DatasetSplit helios.data.datamodule.DataLoaderParams helios.data.datamodule.Dataset helios.data.datamodule.DataModule Functions --------- .. autoapisummary:: helios.data.datamodule.create_dataset helios.data.datamodule.create_collate_fn helios.data.datamodule.create_dataloader Module Contents --------------- .. py:data:: DATASET_REGISTRY Global instance of the registry for datasets. .. rubric:: Example .. code-block:: python 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) .. py:data:: COLLATE_FN_REGISTRY Global instance of the registry for collate functions. .. rubric:: Example .. code-block:: python 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) .. py:function:: 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. :param type_name: the type of the dataset to create. :param args: positional arguments to pass into the dataset. :param kwargs: keyword arguments to pass into the dataset. :returns: The constructed dataset. .. py:function:: 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: #. If ``type_name`` refers to a function-type, then ``args`` and ``kwargs`` are ignored. #. If ``type_name`` refers to an actual type (i.e. a class), then ``args`` and ``kwargs`` are forwarded upon instantiation. :param type_name: the type of the function to create. :param args: positional arguments to pass into the collate class :param kwargs: keyword arguments to pass into the collate class :returns: The function .. py:class:: DatasetSplit Bases: :py:obj:`enum.Enum` The different dataset splits. .. py:attribute:: TRAIN :value: 0 .. py:attribute:: VALID :value: 1 .. py:attribute:: TEST :value: 2 .. py:method:: from_str(label: str) -> DatasetSplit :staticmethod: Convert the given string to the corresponding enum value. Must be one of "train", "test", or "valid" :param label: the label to convert. :returns: The corresponding enum value. :raises ValueError: if the given value is not one of "train", "test", or "valid". .. py:function:: 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 :py:class:`~helios.data.samplers.ResumableDistributedSampler`. * Otherwise, if ``shuffle`` then sampler is :py:class:`~helios.data.samplers.ResumableRandomSampler`, else :py:class:`~helios.data.samplers.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 :py:class:`helios.data.samplers.ResumableSampler` or :py:class:`helios.data.samplers.ResumableDistributedSampler`. :param dataset: the dataset to use. :param random_seed: (optional) seed for the worker processes. Defaults to the value returned by :py:func:`~helios.core.rng.get_default_seed`. :param batch_size: (optional) number of samples per batch. :param shuffle: (optional) if true, samples are randomly shuffled. :param num_workers: (optional) number of worker processes for loading data. :param pin_memory: (optional) if true, use page-locked device memory. :param drop_last: (optional) if true, remove the final batch. :param is_distributed: (optional) if true, create the distributed sampler. :param sampler: (optional) sampler to use. :param collate_fn: (optional) function to merge batches. :param prefetch_factor: (optional) number of batches to prefetch per worker. Only valid when ``num_workers > 0``. :param persistent_workers: (optional) if true, keep worker processes alive between epochs. Only valid when ``num_workers > 0``. :param pin_memory_device: (optional) target device for pinned memory when ``pin_memory`` is true. :param timeout: (optional) timeout in seconds for collecting a batch from workers. :param 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 :py:class:`~helios.data.samplers.ResumableDistributedSampler` or :py:class:`~helios.data.samplers.ResumableSampler`. .. py:class:: DataLoaderParams Params used to create the dataloader object. :param random_seed: (optional) seed for the worker processes. Defaults to the value returned by :py:func:`~helios.core.rng.get_default_seed`. :param batch_size: (optional) number of samples per batch. :param shuffle: (optional) if true, samples are randomly shuffled. :param num_workers: (optional) number of worker processes for loading data. :param pin_memory: (optional) if true, use page-locked device memory. :param drop_last: (optional) if true, remove the final batch. :param is_distributed: (optional) if true, create the distributed sampler. :param sampler: (optional) sampler to use. :param collate_fn: (optional) function to merge batches. :param prefetch_factor: (optional) number of batches to prefetch per worker. Only valid when ``num_workers > 0``. :param persistent_workers: (optional) if true, keep worker processes alive between epochs. Only valid when ``num_workers > 0``. :param pin_memory_device: (optional) target device for pinned memory when ``pin_memory`` is true. :param timeout: (optional) timeout in seconds for collecting a batch from workers. :param multiprocessing_context: (optional) method for spawning worker processes (e.g. ``"fork"``, ``"spawn"``, ``"forkserver"``). .. py:attribute:: random_seed :type: int :value: 6691 .. py:attribute:: batch_size :type: int :value: 1 .. py:attribute:: shuffle :type: bool :value: False .. py:attribute:: num_workers :type: int :value: 0 .. py:attribute:: pin_memory :type: bool :value: False .. py:attribute:: drop_last :type: bool :value: False .. py:attribute:: is_distributed :type: bool | None :value: None .. py:attribute:: sampler :type: helios.data.samplers.ResumableSamplerType | None :value: None .. py:attribute:: collate_fn :type: Callable | None :value: None .. py:attribute:: prefetch_factor :type: int | None :value: None .. py:attribute:: persistent_workers :type: bool :value: False .. py:attribute:: pin_memory_device :type: str :value: '' .. py:attribute:: timeout :type: float :value: 0 .. py:attribute:: multiprocessing_context :type: Any | None :value: None .. py:method:: to_dict() -> dict[str, Any] Convert the params object to a dictionary using shallow copies. .. py:method:: from_dict(table: dict[str, Any]) :classmethod: Create a new params object from the given table. .. py:class:: Dataset The dataset and corresponding data loader params. :param dataset: the dataset. :param params: the data loader params. .. py:attribute:: dataset :type: torch.utils.data.Dataset .. py:attribute:: params :type: DataLoaderParams .. py:method:: dict() -> dict[str, Any] Convert to a dictionary. .. py:class:: DataModule Bases: :py:obj:`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. .. rubric:: Example .. code-block:: python 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. .. py:property:: is_distributed :type: bool Flag controlling whether distributed training is being used or not. .. py:property:: trainer :type: helios.trainer.Trainer Reference to the trainer. .. py:property:: train_dataset :type: torch.utils.data.Dataset | None The training dataset (if available). .. py:property:: valid_dataset :type: torch.utils.data.Dataset | None The validation dataset (if available). .. py:property:: test_dataset :type: torch.utils.data.Dataset | None The testing dataset (if available). .. py:method:: 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. .. py:method:: setup() -> None :abstractmethod: Construct all required datasets. .. py:method:: train_dataloader() -> tuple[torch.utils.data.DataLoader, helios.data.samplers.ResumableSamplerType] | None Create the train dataloader (if available). .. py:method:: valid_dataloader() -> tuple[torch.utils.data.DataLoader, helios.data.samplers.ResumableSamplerType] | None Create the valid dataloader (if available). .. py:method:: test_dataloader() -> tuple[torch.utils.data.DataLoader, helios.data.samplers.ResumableSamplerType] | None Create the test dataloader (if available). .. py:method:: teardown() -> None Clean up any state after training is over. .. py:method:: 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 :py:meth:`~helios.model.model.Model.steup` or :py:meth:`~helios.plugins.plugin.Plugin.setup` and cache the result. This function is equivalent to: .. code-block:: python 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. .. py:method:: 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. .. py:method:: 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. .. py:method:: load_state_dict(state_dict: dict[str, Any]) -> None Load the datamodule state from the given state dictionary. :param state_dict: the state dictionary to load from.