cerebras.modelzoo.data.common.HDF5IterableDataset.HDF5IterableDataset#

class cerebras.modelzoo.data.common.HDF5IterableDataset.HDF5IterableDataset(*args, **kwargs)[source]#

Bases: torch.utils.data.IterableDataset

A HDF5 dataset processor. Loads data from HDF5 files. :param dict params: dict containing training

input parameters for creating dataset.

Expects the following fields: - “data_dir” (str or list of str): Path to dataset HDF5 files - “batch_size” (int): Batch size. - “shuffle” (bool): Flag to enable data shuffling. - “shuffle_seed” (int): Shuffle seed. - “num_workers” (int): How many subprocesses to use for data loading. - “drop_last” (bool): If True and the dataset size is not divisible

by the batch size, the last incomplete batch will be dropped.

  • “use_vsl” (bool): Flag to enable variable sequence length training.

    It requires the dataset to have two extra features: the attention_span of keys and the position_ids of tokens. Defaults to False.

Methods

set_state

This method sets the state of the dataloader's samples_seen variable that controls how many samples are to be skipped for determinisitic restart.

Attributes

samples_seen

set_state(samples_seen, shard_index)[source]#

This method sets the state of the dataloader’s samples_seen variable that controls how many samples are to be skipped for determinisitic restart. This is called by the load_state_dict method of the RestartableDataLoader.

Parameters
  • samples_seen (int) – number of samples streamed by the dataloader

  • shard_index (int) – the index of the shard of data that this worker is responsible for streaming