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
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