Class DataSpec

composer/core/data_spec.py:126–318

Specifications for operating and training on data. An example of constructing a :class:`DataSpec` object with a ``batch_transforms`` callable and then using it with :class:`~.Trainer`: .. doctest:: >>> # Construct DataSpec and subtract mean from the batch >>> batch_trans

class DataSpec:
    """Specifications for operating and training on data.

    An example of constructing a :class:`DataSpec` object with a ``batch_transforms``
    callable and then using it with :class:`~.Trainer`:

    .. doctest::

       >>> # Construct DataSpec and subtract mean from the batch
       >>> batch_transform_fn = lambda xs, ys: (xs.sub_(xs.mean()), ys)
       >>> train_dspec = DataSpec(train_dataloader, batch_transforms=batch_transform_fn)
       >>> # The same function can be used for eval dataloader as well
       >>> eval_dspec = DataSpec(eval_dataloader, batch_transforms=batch_transform_fn)
       >>> # Use this DataSpec object to construct trainer
       >>> trainer = Trainer(
       ...     model=model,
       ...     train_dataloader=train_dspec,
       ...     eval_dataloader=eval_dspec,
       ...     optimizers=optimizer,
       ...     max_duration="1ep",
       ... )

    Args:
        dataloader (Union[Iterable, torch.utils.data.DataLoader]): The dataloader, which can be any iterable that yields batches.

        num_samples (int, optional): The total number of samples in an epoch, across all ranks. This field is used by
            the :class:`.Timestamp` (training progress tracker). If not specified, then ``len(dataloader.dataset)`` is
            used (if this property is available). Otherwise, the dataset is assumed to be unsized.

        num_tokens (int, optional): The total number of tokens in an epoch. This field is used by the
            :class:`.Timestamp` (training progress tracker).

        batch_transforms ((Batch) -> Batch, optional): Function called by the :class:`.Trainer` to modify the
            batch before it is moved onto the device. For example, this function can be used for CPU-based
            normalization. It can modify the batch in-place, and it should return the modified batch. If not specified,
            the batch is not modified.

        microbatch_transforms ((Batch) -> Batch, optional): Function called by the :class:`.Trainer` to modify the
            microbatch before it is moved onto the device. For example, this function can be used for GPU-based
            normalization. It can modify the microbatch in-place, and it should return the modified microbatch. If not
            specified, the microbatch is not modified.

        split_batch ((Batch, (int | float)) -> Sequence[Batch], optional): Function called by the :class:`.Trainer` to
            split a batch (the first parameter) into microbatches of a given size (the second parameter). If
            the ``dataloader`` yields batches not of type :class:`torch.Tensor`, Mapping, tuple, or list, then
            this function must be specified.

        get_num_samples_in_batch ((Batch) -> Union[int, float], optional): Function that is called by the :class:`.Trainer`
            to get the number of samples in the provided batch.

            By default, if the batch contains tensors that all have the same 0th dim, then the value of the 0th dim will
            be returned. If the batch contains tensors where the 0th dim differ, then this function must be specified.

        get_num_tokens_in_batch ((Batch) -> int, optional): Function that is called by the :class:`.Trainer` to
            get the number of tokens in the provided batch.

            By default, it checks for HuggingFace-style dictionary batches with ``input_ids``, and then checks ``dataset.max_seq_len``, and returns 0
            if both of those fail, meaning that number of tokens processed will not be tracked as a part of the training progress tracking.