Class BatchPreparer

cdslib/core/data/dataloader/batch_preparer.py:26–287

The batch preparer works as follows: - First adds datasets before iterator starts. - When creating the iterator (beginning of a epoch): 1. It determines the samples to used in the iteration based on dataset_ratios. 2. It groups samples of similar length into batch (usin

class BatchPreparer:
    """
    The batch preparer works as follows:

    - First adds datasets before iterator starts.
    - When creating the iterator (beginning of a epoch):
        1. It determines the samples to used in the iteration based on dataset_ratios.
        2. It groups samples of similar length into batch (using bucket_sampler or sort_sampler).
    - Every iteration:
        The samples are loaded in parallel using pytorch's dataloader.

    Data flow:

    - Each dataset is stored in a pytorch's Dataset (with an additional method 'get_seq_lens' implemented).

        ex: seq_lens = dataset.get_seq_lens()

        - seq_lens is a list of all samples stored in the dataset.
        - seq_lens[i] = length of dataset[i]

    - The datasets are concatenated to a single dataset by torch.utils.data.ConcatDataset for pytorch's dataloader.
    """

    def __init__(
        self,
        batch_size,
        collate_fn,
        drop_last=False,
        num_workers=0,
        batch_wrt_length=True,
        shuffle=True,
        max_total_samples_per_epoch=-1,
        max_batch_combined_size=-1,
        batch_sampler_type="sort",
    ):
        """
        Args:
            batch_size:
                typical batch_size
            collate_fn:
                collate function used to combine samples from dataset
            drop_last:
                whether to drop the rest of the data that cannot form a batch
            num_workers:
                number of loading threads to use
            batch_wrt_length:
                whether to batch samples of similar sequence lengths
            shuffle:
                whether to shuffle data within each bucket
            max_total_samples_per_epoch:
                maximum total number of samples per epoch. -1: ignored
            max_batch_combined_size:
                limitation on batch_size * seq_len.  -1: ignored
            batch_sampler_type:
                type of the batch sampler [bucket | sort]
        """

        self.batch_size = batch_size