Class BucketSampler

cdslib/core/data/dataloader/bucket_sampler.py:16–215

The class assembles batches by grouping data of similar length. It returns indices of the data in the dataset within a batch. Here is an example to use the bucket sampler: .. code-block:: python from bucket_sampler import Bucket_Sampler import numpy as np

class BucketSampler:
    """
    The class assembles batches by grouping data of similar length.
    It returns indices of the data in the dataset within a batch.

    Here is an example to use the bucket sampler:

    .. code-block:: python

        from bucket_sampler import Bucket_Sampler
        import numpy as np

        dataset_size = 50
        dataset_lengths = np.random.randint(1, 100, dataset_size)
        print(dataset_lengths)

        batch_size = 4
        drop_last = True
        shuffle = True
        bucket_boundaries = [10,20,30,40,50,60,70,80,90]  # I hand-assigned here, but it can be created by np.histogram_bin_edges
        batch_sampler = Bucket_Sampler(dataset_lengths, bucket_boundaries, batch_size, drop_last, shuffle)

        print('num batches = %d' % (len(batch_sampler)))
        for batch_idx, data_idxs in enumerate(batch_sampler):
            print('%d: '% batch_idx, end='' )
            print(data_idxs)

        # no need to assign batch_size, shuffle, drop_last, since batch_sampler determines them
        dataloader = torch.utils.data.DataLoader(
            self.dataset,
            batch_sampler=bucket_sampler,
            num_workers=self.num_workers,
            collate_fn=self.dataset.collate_fn
        )

    """

    def __init__(
        self,
        batch_size,
        drop_last,
        shuffle,
        bucket_boundaries=None,
        max_batch_combined_size=-1,
    ):
        """
        Args:
            batch_size:
                batch size
            drop_last:
                whether to drop the last few data that cannot form a batch (recommanded True).
            shuffle:
                whether to shuffle the data slightly (according to their length of course).
            bucket_boundaries:
                int (number of bins) or a list (containing the edges in ascending order,
                excluding two outmost boundaries).

                edge     0   1   2