Sampler

BucketBatchSampler

Bases: Sampler[List[int]]

A batch sampler to create batches with sizes of elements from each pre-defined bucket ranges.

Elements of the dataset are first grouped into each bucket based on the bucket ranges and the sizes of elements. Then, a base batch sampler is used for each bucket to create mini-batches.

The bucket ranges are specified by bucket_boundaries, which will be first sorted internally and used to create len(bucket_boundaries) - 1 left-closed right-open intervals. e.g. if bucket_boundaries tensor is [10, 5, 0, 16], it will be sorted as [0, 5, 10, 16] and 3 buckets will be created with ranges: [0, 5), [5, 10), [10, 16).

The base batch sampler will be created by passing the element indices in each bucket as the data source, and base_batch_sampler_shared_kwargs and base_batch_sampler_individual_kwargs to the constructor of the base batch sampler class specified as base_batch_sampler_class. e.g. base_batch_sampler_shared_kwargs = {'drop_last': True} and base_batch_sampler_individual_kwargs = {'batch_size': [8,10,12]} will be used to create 3 batch samplers with drop_last=True and batch_size=8, 10 and 12, and initialized like base_batch_sampler_class(bucket_element_indices[0], batch_size=8, drop_last=True).

In the __iter__ method, if shuffle is True, the element indices in each bucket will be shuffled, and a bucket is randomly selected each time to create a mini-batch. If shuffle is False, there is no shuffle on element indices, and the bucket is selected in ascending order of its interval boundaries.

This class is used to create homogeneous batches of data for training or evaluation, and reduce the padding necessary to align the shape of elements.

Modified from https://github.com/rssrwn/semla-flow/blob/main/semlaflow/data/util.py

---

Examples:

>>> import torch
>>> from bionemo.size_aware_batching.sampler import BucketBatchSampler

>>> # Define the sizes for a dataset
>>> sizes = torch.arange(25)
>>> # Define bucket ranges
>>> bucket_boundaries = torch.tensor([0, 6, 15, 25])

>>> # Create a bucket batch sampler with torch.utils.data.BatchSampler as base batch sampler
>>> # As there are 3 buckets, there will be 3 base batch samplers with batch sizes 2, 3, and 5.
>>> batch_sampler = BucketBatchSampler(
        sizes=sizes,
        bucket_boundaries=bucket_boundaries,
        base_batch_sampler_class=torch.utils.data.BatchSampler,
        base_batch_sampler_shared_kwargs={'drop_last': False},
        base_batch_sampler_individual_kwargs={'batch_size': [2,3,5]},
        shuffle=False,
    )

>>> # Iterate over batches of indices that lies in the same bucket and with different batch sizes.
>>> print(list(batch_sampler))
[[0, 1], [2, 3], [4, 5], [6, 7, 8], [9, 10, 11], [12, 13, 14], [15, 16, 17, 18, 19], [20, 21, 22, 23, 24]]

>>> # randomize the dataset and buckets
>>> batch_sampler = BucketBatchSampler(
        sizes=sizes,
        bucket_boundaries=bucket_boundaries,
        base_batch_sampler_class=torch.utils.data.BatchSampler,
        base_batch_sampler_shared_kwargs={'drop_last': False},
        base_batch_sampler_individual_kwargs={'batch_size': [2,3,5]},
        shuffle=True,
        generator=torch.Generator().manual_seed(0),
    )
>>> print(list(batch_sampler))
[[24, 17, 16, 22, 19], [2, 5], [12, 10, 11], [3, 0], [15, 18, 20, 21, 23], [7, 13, 6], [14, 9, 8], [1, 4]]
>>> print(list(batch_sampler))
[[14, 9, 13], [23, 16, 20, 21, 15], [5, 0], [8, 10, 11], [17, 24, 22, 18, 19], [12, 6, 7], [4, 2], [3, 1]]

>>> # Combine with SizeAwareBatchSampler to control the cost of each batch
>>> from bionemo.size_aware_batching.sampler import SizeAwareBatchSampler
>>> item_costs = sizes.tolist()
>>> def cost_of_element(index):
        return item_costs[index]
>>> batch_sampler = BucketBatchSampler(
        sizes=sizes,
        bucket_boundaries=bucket_boundaries,
        base_batch_sampler_class=SizeAwareBatchSampler,
        base_batch_sampler_shared_kwargs={"sizeof": cost_of_element, "max_total_size": 40},
        base_batch_sampler_individual_kwargs={},
        shuffle=True,
        generator=torch.Generator().manual_seed(0),
    )
>>> print(list(iter(batch_sampler)))
[[24], [2, 5, 3, 0, 1, 4], [12, 10, 11, 7], [13, 6, 14], [17, 16], [22], [19, 15], [9, 8], [18, 20], [21], [23]]

Source code in bionemo/size_aware_batching/sampler.py

class BucketBatchSampler(Sampler[List[int]]):
    """A batch sampler to create batches with sizes of elements from each pre-defined bucket ranges.

    Elements of the dataset are first grouped into each bucket based on the bucket ranges and the sizes of elements.
    Then, a base batch sampler is used for each bucket to create mini-batches.

    The bucket ranges are specified by `bucket_boundaries`, which will be first sorted internally and used to create
    `len(bucket_boundaries) - 1` left-closed right-open intervals.
    e.g. if bucket_boundaries tensor is [10, 5, 0, 16], it will be sorted as [0, 5, 10, 16] and 3 buckets will be created
    with ranges: [0, 5), [5, 10), [10, 16).

    The base batch sampler will be created by passing the element indices in each bucket as the data source, and
    `base_batch_sampler_shared_kwargs` and `base_batch_sampler_individual_kwargs`
    to the constructor of the base batch sampler class specified as `base_batch_sampler_class`.
    e.g. `base_batch_sampler_shared_kwargs = {'drop_last': True}` and `base_batch_sampler_individual_kwargs = {'batch_size': [8,10,12]}`
    will be used to create 3 batch samplers with drop_last=True and batch_size=8, 10 and 12, and initialized like
    `base_batch_sampler_class(bucket_element_indices[0], batch_size=8, drop_last=True)`.

    In the `__iter__` method, if `shuffle` is `True`, the element indices in each bucket will be shuffled, and a bucket
    is randomly selected each time to create a mini-batch. If `shuffle` is `False`, there is no shuffle on element indices,
    and the bucket is selected in ascending order of its interval boundaries.

    This class is used to create homogeneous batches of data for training or evaluation, and reduce the padding necessary to align the shape of elements.

    Modified from https://github.com/rssrwn/semla-flow/blob/main/semlaflow/data/util.py

    ---------

    Examples:
    ```python
    >>> import torch
    >>> from bionemo.size_aware_batching.sampler import BucketBatchSampler

    >>> # Define the sizes for a dataset
    >>> sizes = torch.arange(25)
    >>> # Define bucket ranges
    >>> bucket_boundaries = torch.tensor([0, 6, 15, 25])

    >>> # Create a bucket batch sampler with torch.utils.data.BatchSampler as base batch sampler
    >>> # As there are 3 buckets, there will be 3 base batch samplers with batch sizes 2, 3, and 5.
    >>> batch_sampler = BucketBatchSampler(
            sizes=sizes,
            bucket_boundaries=bucket_boundaries,
            base_batch_sampler_class=torch.utils.data.BatchSampler,
            base_batch_sampler_shared_kwargs={'drop_last': False},
            base_batch_sampler_individual_kwargs={'batch_size': [2,3,5]},
            shuffle=False,
        )

    >>> # Iterate over batches of indices that lies in the same bucket and with different batch sizes.
    >>> print(list(batch_sampler))
    [[0, 1], [2, 3], [4, 5], [6, 7, 8], [9, 10, 11], [12, 13, 14], [15, 16, 17, 18, 19], [20, 21, 22, 23, 24]]

    >>> # randomize the dataset and buckets
    >>> batch_sampler = BucketBatchSampler(
            sizes=sizes,
            bucket_boundaries=bucket_boundaries,
            base_batch_sampler_class=torch.utils.data.BatchSampler,
            base_batch_sampler_shared_kwargs={'drop_last': False},
            base_batch_sampler_individual_kwargs={'batch_size': [2,3,5]},
            shuffle=True,
            generator=torch.Generator().manual_seed(0),
        )
    >>> print(list(batch_sampler))
    [[24, 17, 16, 22, 19], [2, 5], [12, 10, 11], [3, 0], [15, 18, 20, 21, 23], [7, 13, 6], [14, 9, 8], [1, 4]]
    >>> print(list(batch_sampler))
    [[14, 9, 13], [23, 16, 20, 21, 15], [5, 0], [8, 10, 11], [17, 24, 22, 18, 19], [12, 6, 7], [4, 2], [3, 1]]

    >>> # Combine with SizeAwareBatchSampler to control the cost of each batch
    >>> from bionemo.size_aware_batching.sampler import SizeAwareBatchSampler
    >>> item_costs = sizes.tolist()
    >>> def cost_of_element(index):
            return item_costs[index]
    >>> batch_sampler = BucketBatchSampler(
            sizes=sizes,
            bucket_boundaries=bucket_boundaries,
            base_batch_sampler_class=SizeAwareBatchSampler,
            base_batch_sampler_shared_kwargs={"sizeof": cost_of_element, "max_total_size": 40},
            base_batch_sampler_individual_kwargs={},
            shuffle=True,
            generator=torch.Generator().manual_seed(0),
        )
    >>> print(list(iter(batch_sampler)))
    [[24], [2, 5, 3, 0, 1, 4], [12, 10, 11, 7], [13, 6, 14], [17, 16], [22], [19, 15], [9, 8], [18, 20], [21], [23]]
    ```
    """

    def __init__(
        self,
        sizes: torch.Tensor,
        bucket_boundaries: torch.Tensor,
        base_batch_sampler_class: Type[S],
        base_batch_sampler_shared_kwargs: Optional[Dict[str, Any]] = None,
        base_batch_sampler_individual_kwargs: Optional[Dict[str, Iterable]] = None,
        shuffle: Optional[bool] = True,
        generator: Optional[torch.Generator] = None,
    ) -> None:
        """Initializes the BucketBatchSampler.

        Args:
            sizes: A 1D tensor of real numbers representing the size of each element in the dataset.
            bucket_boundaries: A 1D tensor of real numbers representing the boundaries of the bucket ranges.
                It will be first sorted and used to create `len(bucket_boundaries) - 1` left-closed right-open intervals as bucket ranges.
                It should not contain any duplicate values.
            base_batch_sampler_class: Base batch sampler class type, which will be used for each bucket, and initialized with the bucket element indices,
                `base_batch_sampler_shared_kwargs` and the corresponding `base_batch_sampler_individual_kwargs`.
            base_batch_sampler_shared_kwargs: Shared keyword argument dictionary used to initialize all base batch samplers for all buckets.
                Sufficient and valid arguments should be provided for `base_batch_sampler_class` with `base_batch_sampler_individual_kwargs`. Default to  {}.
            base_batch_sampler_individual_kwargs: Keyword argument dictionary used to initialize
                each bucket batch sampler with the corresponding key value pairs.
                Length of each value in this dict must be equal to len(bucket_boundaries) - 1 (the number of buckets).
                Sufficient and valid arguments should be provided for `base_batch_sampler_class` with `base_batch_sampler_shared_kwargs`.
                Default to  {}.
            shuffle: A boolean indicating whether to shuffle the dataset and buckets. Defaults to True.
            generator: Generator used in sampling. Defaults to None.

        Raises:
            ValueError: If `sizes` is not a 1D tensor of real numbers.
            ValueError: If `bucket_boundaries` is not a 1D tensor of real numbers.
            ValueError: If `base_batch_sampler_individual_kwargs` or `base_batch_sampler_individual_kwargs` is not a keyword argument dictionary.
            ValueError: If the length of values in the dict of `base_batch_sampler_individual_kwargs` must be equal to len(bucket_boundaries) - 1.
            RuntimeError: If there is no elements with sizes inside the ranges specified by `bucket_boundaries`.

        """
        if not torch.is_tensor(sizes):
            raise TypeError(f"sizes should be a torch tensor, but got sizes={sizes}")

        if sizes.ndim != 1:
            raise ValueError(f"sizes should be a 1D tensor, but got sizes with shape {sizes.shape}")

        if not torch.is_floating_point(sizes) and sizes.dtype not in TorchIntegerDataTypes:
            raise ValueError(
                f"sizes should contain only integers or floating point numbers, but got sizes.dtype={sizes.dtype}"
            )

        if not torch.is_tensor(bucket_boundaries):
            raise TypeError(
                f"bucket_boundaries should be a torch tensor, but got bucket_boundaries={bucket_boundaries}"
            )

        if bucket_boundaries.ndim != 1:
            raise ValueError(
                f"bucket_boundaries should be a 2D tensor, but got bucket_boundaries with shape {bucket_boundaries.shape}"
            )

        if len(bucket_boundaries) < 2:
            raise ValueError(
                f"bucket_boundaries should have at least 2 numbers, but got bucket_boundaries={bucket_boundaries.shape}"
            )

        if not torch.is_floating_point(bucket_boundaries) and bucket_boundaries.dtype not in TorchIntegerDataTypes:
            raise ValueError(
                f"bucket_boundaries should contain only integers or floating point numbers, but got bucket_boundaries.dtype={bucket_boundaries.dtype}"
            )

        bucket_boundaries = torch.sort(bucket_boundaries)[0]

        if torch.any(bucket_boundaries[:-1] >= bucket_boundaries[1:]):
            raise ValueError(
                f"bucket_boundaries should not have duplicate values, and should specify the lower endpoint of each interval smaller than the upper endpoint, but got sorted bucket_boundaries={bucket_boundaries}"
            )

        if not isinstance(shuffle, bool):
            raise TypeError(f"shuffle should be a boolean value, but got shuffle={shuffle}")

        self.sizes = sizes
        self.bucket_boundaries = bucket_boundaries
        self.num_buckets = len(bucket_boundaries) - 1
        self.shuffle = shuffle
        self.generator = generator
        if self.shuffle and self.generator is None:
            self.generator = torch.Generator().manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))

        if not issubclass(base_batch_sampler_class, Sampler):
            raise TypeError(
                f"base_batch_sampler_class should be a batch sampler class inherited from torch.utils.data.Sampler, but got base_batch_sampler_class={base_batch_sampler_class}"
            )

        if not isinstance(base_batch_sampler_shared_kwargs, dict):
            raise TypeError(
                f"base_batch_sampler_shared_kwargs should be a dictionary, but got base_batch_sampler_shared_kwargs={base_batch_sampler_shared_kwargs}"
            )

        if not all(isinstance(key, str) for key in base_batch_sampler_shared_kwargs.keys()):
            raise TypeError(
                f"base_batch_sampler_shared_kwargs should have string keys, but got keys={list(base_batch_sampler_shared_kwargs.keys())}"
            )

        if not isinstance(base_batch_sampler_individual_kwargs, dict):
            raise TypeError(
                f"base_batch_sampler_individual_kwargs should be a dictionary, but got base_batch_sampler_individual_kwargs={base_batch_sampler_individual_kwargs}"
            )

        if not all(isinstance(key, str) for key in base_batch_sampler_individual_kwargs.keys()):
            raise TypeError(
                f"base_batch_sampler_individual_kwargs should have string keys, but got keys={list(base_batch_sampler_individual_kwargs.keys())}"
            )

        if not all(len(list(value)) == self.num_buckets for value in base_batch_sampler_individual_kwargs.values()):
            raise ValueError(
                f"Each value in base_batch_sampler_individual_kwargs should have a length of {self.num_buckets}, "
                f"but got lengths {[len(list(value)) for value in base_batch_sampler_individual_kwargs.values()]}"
            )

        self.base_batch_sampler_class = base_batch_sampler_class
        self.base_batch_sampler_shared_kwargs = (
            {} if base_batch_sampler_shared_kwargs is None else base_batch_sampler_shared_kwargs
        )
        base_batch_sampler_individual_kwargs = (
            {} if base_batch_sampler_individual_kwargs is None else base_batch_sampler_individual_kwargs
        )
        self.base_batch_sampler_individual_kwargs = [
            {key: list(base_batch_sampler_individual_kwargs[key])[k] for key in base_batch_sampler_individual_kwargs}
            for k in range(self.num_buckets)
        ]

        self.bucket_sizes: torch.Tensor  # number of elements in each bucket
        self.bucket_element_indices: List[List[int]]  # List of elements' indices for each bucket

        # bucket index for each element
        element_bucket_indices = torch.bucketize(sizes, bucket_boundaries, right=True)

        # element indices reordered for each bucket
        reordered_element_indices = torch.argsort(element_bucket_indices, stable=True)

        # bucket sizes, including the buckets for < bucket_boundaries[0] and >= bucket_boundaries[-1]
        bucket_sizes = torch.bincount(element_bucket_indices, minlength=len(bucket_boundaries) + 1)

        # bucket segments
        bucket_segments = torch.cumsum(bucket_sizes, dim=0)[:-1]

        self.bucket_element_indices = []
        # exclude the buckets for < bucket_boundaries[0] and >= bucket_boundaries[-1]
        for bucket_idx in range(self.num_buckets):
            self.bucket_element_indices.append(
                reordered_element_indices[bucket_segments[bucket_idx] : bucket_segments[bucket_idx + 1]].tolist()
            )
        self.bucket_sizes = bucket_sizes[1 : (self.num_buckets + 1)]

        self.num_samples = torch.sum(self.bucket_sizes).item()
        if self.num_samples == 0:
            raise RuntimeError("The sizes of all elements in the dataset are outside the bucket ranges provided")
        if self.num_samples < len(self.sizes):
            warnings.warn(
                f"{len(self.sizes) - self.num_samples} elements are outside the buckets provided and will be skipped"
            )

        self.base_batch_samplers: List[Sampler] = self._init_base_batch_samplers()

    def _init_base_batch_samplers(self) -> list[Sampler[List[int]]]:
        """Initialize batch samplers for each bucket.

        Returns:
            List of batch samplers.
        """
        base_batch_samplers = []
        for k in range(self.num_buckets):
            base_batch_samplers.append(
                self.base_batch_sampler_class(
                    self.bucket_element_indices[k],
                    **self.base_batch_sampler_shared_kwargs,
                    **self.base_batch_sampler_individual_kwargs[k],
                )
            )
        return base_batch_samplers

    def __len__(self) -> int:
        """Get the number of batches.

        Can only be called if the `base_batch_sampler_class` has __len__() implemented

        Returns:
            int: Number of batches
        """
        num_batches = sum(len(sampler) for sampler in self.base_batch_samplers)  # type: ignore
        return num_batches

    def __iter__(self) -> Iterator[List[int]]:
        """Iterate over batches of indices.

        This function yields batches of indices of elements with sizes from each bucket range.

        Yields:
            List[int]: A batch of indices of elements with sizes from each bucket range.
        """
        if self.shuffle:
            for indices in self.bucket_element_indices:
                idx = torch.randperm(len(indices), generator=self.generator)
                indices[:] = torch.tensor(indices)[idx].tolist()

        base_batch_sampler_iters = [iter(batch_sampler) for batch_sampler in self.base_batch_samplers]
        bucket_remaining_elements = self.bucket_sizes.clone()
        total_remaining_elements = self.num_samples

        while total_remaining_elements > 0:
            if self.shuffle:
                bucket_idx = torch.multinomial(
                    bucket_remaining_elements / total_remaining_elements, 1, generator=self.generator
                )
            else:
                bucket_idx = torch.argmax((bucket_remaining_elements > 0).to(int))  # type: ignore

            try:
                batch = next(base_batch_sampler_iters[bucket_idx])
                bucket_remaining_elements[bucket_idx] -= len(batch)
                total_remaining_elements -= len(batch)
                yield batch
            except StopIteration:
                bucket_remaining_elements[bucket_idx] = 0
                total_remaining_elements = torch.sum(bucket_remaining_elements)
                continue

__init__(sizes, bucket_boundaries, base_batch_sampler_class, base_batch_sampler_shared_kwargs=None, base_batch_sampler_individual_kwargs=None, shuffle=True, generator=None)

Initializes the BucketBatchSampler.

Parameters:

Name	Type	Description	Default
sizes	Tensor	A 1D tensor of real numbers representing the size of each element in the dataset.	required
bucket_boundaries	Tensor	A 1D tensor of real numbers representing the boundaries of the bucket ranges. It will be first sorted and used to create len(bucket_boundaries) - 1 left-closed right-open intervals as bucket ranges. It should not contain any duplicate values.	required
base_batch_sampler_class	Type[S]	Base batch sampler class type, which will be used for each bucket, and initialized with the bucket element indices, base_batch_sampler_shared_kwargs and the corresponding base_batch_sampler_individual_kwargs.	required
base_batch_sampler_shared_kwargs	Optional[Dict[str, Any]]	Shared keyword argument dictionary used to initialize all base batch samplers for all buckets. Sufficient and valid arguments should be provided for base_batch_sampler_class with base_batch_sampler_individual_kwargs. Default to {}.	None
base_batch_sampler_individual_kwargs	Optional[Dict[str, Iterable]]	Keyword argument dictionary used to initialize each bucket batch sampler with the corresponding key value pairs. Length of each value in this dict must be equal to len(bucket_boundaries) - 1 (the number of buckets). Sufficient and valid arguments should be provided for base_batch_sampler_class with base_batch_sampler_shared_kwargs. Default to {}.	None
shuffle	Optional[bool]	A boolean indicating whether to shuffle the dataset and buckets. Defaults to True.	True
generator	Optional[Generator]	Generator used in sampling. Defaults to None.	None

Raises:

Type	Description
ValueError	If sizes is not a 1D tensor of real numbers.
ValueError	If bucket_boundaries is not a 1D tensor of real numbers.
ValueError	If base_batch_sampler_individual_kwargs or base_batch_sampler_individual_kwargs is not a keyword argument dictionary.
ValueError	If the length of values in the dict of base_batch_sampler_individual_kwargs must be equal to len(bucket_boundaries) - 1.
RuntimeError	If there is no elements with sizes inside the ranges specified by bucket_boundaries.

Source code in bionemo/size_aware_batching/sampler.py

def __init__(
    self,
    sizes: torch.Tensor,
    bucket_boundaries: torch.Tensor,
    base_batch_sampler_class: Type[S],
    base_batch_sampler_shared_kwargs: Optional[Dict[str, Any]] = None,
    base_batch_sampler_individual_kwargs: Optional[Dict[str, Iterable]] = None,
    shuffle: Optional[bool] = True,
    generator: Optional[torch.Generator] = None,
) -> None:
    """Initializes the BucketBatchSampler.

    Args:
        sizes: A 1D tensor of real numbers representing the size of each element in the dataset.
        bucket_boundaries: A 1D tensor of real numbers representing the boundaries of the bucket ranges.
            It will be first sorted and used to create `len(bucket_boundaries) - 1` left-closed right-open intervals as bucket ranges.
            It should not contain any duplicate values.
        base_batch_sampler_class: Base batch sampler class type, which will be used for each bucket, and initialized with the bucket element indices,
            `base_batch_sampler_shared_kwargs` and the corresponding `base_batch_sampler_individual_kwargs`.
        base_batch_sampler_shared_kwargs: Shared keyword argument dictionary used to initialize all base batch samplers for all buckets.
            Sufficient and valid arguments should be provided for `base_batch_sampler_class` with `base_batch_sampler_individual_kwargs`. Default to  {}.
        base_batch_sampler_individual_kwargs: Keyword argument dictionary used to initialize
            each bucket batch sampler with the corresponding key value pairs.
            Length of each value in this dict must be equal to len(bucket_boundaries) - 1 (the number of buckets).
            Sufficient and valid arguments should be provided for `base_batch_sampler_class` with `base_batch_sampler_shared_kwargs`.
            Default to  {}.
        shuffle: A boolean indicating whether to shuffle the dataset and buckets. Defaults to True.
        generator: Generator used in sampling. Defaults to None.

    Raises:
        ValueError: If `sizes` is not a 1D tensor of real numbers.
        ValueError: If `bucket_boundaries` is not a 1D tensor of real numbers.
        ValueError: If `base_batch_sampler_individual_kwargs` or `base_batch_sampler_individual_kwargs` is not a keyword argument dictionary.
        ValueError: If the length of values in the dict of `base_batch_sampler_individual_kwargs` must be equal to len(bucket_boundaries) - 1.
        RuntimeError: If there is no elements with sizes inside the ranges specified by `bucket_boundaries`.

    """
    if not torch.is_tensor(sizes):
        raise TypeError(f"sizes should be a torch tensor, but got sizes={sizes}")

    if sizes.ndim != 1:
        raise ValueError(f"sizes should be a 1D tensor, but got sizes with shape {sizes.shape}")

    if not torch.is_floating_point(sizes) and sizes.dtype not in TorchIntegerDataTypes:
        raise ValueError(
            f"sizes should contain only integers or floating point numbers, but got sizes.dtype={sizes.dtype}"
        )

    if not torch.is_tensor(bucket_boundaries):
        raise TypeError(
            f"bucket_boundaries should be a torch tensor, but got bucket_boundaries={bucket_boundaries}"
        )

    if bucket_boundaries.ndim != 1:
        raise ValueError(
            f"bucket_boundaries should be a 2D tensor, but got bucket_boundaries with shape {bucket_boundaries.shape}"
        )

    if len(bucket_boundaries) < 2:
        raise ValueError(
            f"bucket_boundaries should have at least 2 numbers, but got bucket_boundaries={bucket_boundaries.shape}"
        )

    if not torch.is_floating_point(bucket_boundaries) and bucket_boundaries.dtype not in TorchIntegerDataTypes:
        raise ValueError(
            f"bucket_boundaries should contain only integers or floating point numbers, but got bucket_boundaries.dtype={bucket_boundaries.dtype}"
        )

    bucket_boundaries = torch.sort(bucket_boundaries)[0]

    if torch.any(bucket_boundaries[:-1] >= bucket_boundaries[1:]):
        raise ValueError(
            f"bucket_boundaries should not have duplicate values, and should specify the lower endpoint of each interval smaller than the upper endpoint, but got sorted bucket_boundaries={bucket_boundaries}"
        )

    if not isinstance(shuffle, bool):
        raise TypeError(f"shuffle should be a boolean value, but got shuffle={shuffle}")

    self.sizes = sizes
    self.bucket_boundaries = bucket_boundaries
    self.num_buckets = len(bucket_boundaries) - 1
    self.shuffle = shuffle
    self.generator = generator
    if self.shuffle and self.generator is None:
        self.generator = torch.Generator().manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))

    if not issubclass(base_batch_sampler_class, Sampler):
        raise TypeError(
            f"base_batch_sampler_class should be a batch sampler class inherited from torch.utils.data.Sampler, but got base_batch_sampler_class={base_batch_sampler_class}"
        )

    if not isinstance(base_batch_sampler_shared_kwargs, dict):
        raise TypeError(
            f"base_batch_sampler_shared_kwargs should be a dictionary, but got base_batch_sampler_shared_kwargs={base_batch_sampler_shared_kwargs}"
        )

    if not all(isinstance(key, str) for key in base_batch_sampler_shared_kwargs.keys()):
        raise TypeError(
            f"base_batch_sampler_shared_kwargs should have string keys, but got keys={list(base_batch_sampler_shared_kwargs.keys())}"
        )

    if not isinstance(base_batch_sampler_individual_kwargs, dict):
        raise TypeError(
            f"base_batch_sampler_individual_kwargs should be a dictionary, but got base_batch_sampler_individual_kwargs={base_batch_sampler_individual_kwargs}"
        )

    if not all(isinstance(key, str) for key in base_batch_sampler_individual_kwargs.keys()):
        raise TypeError(
            f"base_batch_sampler_individual_kwargs should have string keys, but got keys={list(base_batch_sampler_individual_kwargs.keys())}"
        )

    if not all(len(list(value)) == self.num_buckets for value in base_batch_sampler_individual_kwargs.values()):
        raise ValueError(
            f"Each value in base_batch_sampler_individual_kwargs should have a length of {self.num_buckets}, "
            f"but got lengths {[len(list(value)) for value in base_batch_sampler_individual_kwargs.values()]}"
        )

    self.base_batch_sampler_class = base_batch_sampler_class
    self.base_batch_sampler_shared_kwargs = (
        {} if base_batch_sampler_shared_kwargs is None else base_batch_sampler_shared_kwargs
    )
    base_batch_sampler_individual_kwargs = (
        {} if base_batch_sampler_individual_kwargs is None else base_batch_sampler_individual_kwargs
    )
    self.base_batch_sampler_individual_kwargs = [
        {key: list(base_batch_sampler_individual_kwargs[key])[k] for key in base_batch_sampler_individual_kwargs}
        for k in range(self.num_buckets)
    ]

    self.bucket_sizes: torch.Tensor  # number of elements in each bucket
    self.bucket_element_indices: List[List[int]]  # List of elements' indices for each bucket

    # bucket index for each element
    element_bucket_indices = torch.bucketize(sizes, bucket_boundaries, right=True)

    # element indices reordered for each bucket
    reordered_element_indices = torch.argsort(element_bucket_indices, stable=True)

    # bucket sizes, including the buckets for < bucket_boundaries[0] and >= bucket_boundaries[-1]
    bucket_sizes = torch.bincount(element_bucket_indices, minlength=len(bucket_boundaries) + 1)

    # bucket segments
    bucket_segments = torch.cumsum(bucket_sizes, dim=0)[:-1]

    self.bucket_element_indices = []
    # exclude the buckets for < bucket_boundaries[0] and >= bucket_boundaries[-1]
    for bucket_idx in range(self.num_buckets):
        self.bucket_element_indices.append(
            reordered_element_indices[bucket_segments[bucket_idx] : bucket_segments[bucket_idx + 1]].tolist()
        )
    self.bucket_sizes = bucket_sizes[1 : (self.num_buckets + 1)]

    self.num_samples = torch.sum(self.bucket_sizes).item()
    if self.num_samples == 0:
        raise RuntimeError("The sizes of all elements in the dataset are outside the bucket ranges provided")
    if self.num_samples < len(self.sizes):
        warnings.warn(
            f"{len(self.sizes) - self.num_samples} elements are outside the buckets provided and will be skipped"
        )

    self.base_batch_samplers: List[Sampler] = self._init_base_batch_samplers()

__iter__()

Iterate over batches of indices.

This function yields batches of indices of elements with sizes from each bucket range.

Yields:

Type	Description
List[int]	List[int]: A batch of indices of elements with sizes from each bucket range.

Source code in bionemo/size_aware_batching/sampler.py

def __iter__(self) -> Iterator[List[int]]:
    """Iterate over batches of indices.

    This function yields batches of indices of elements with sizes from each bucket range.

    Yields:
        List[int]: A batch of indices of elements with sizes from each bucket range.
    """
    if self.shuffle:
        for indices in self.bucket_element_indices:
            idx = torch.randperm(len(indices), generator=self.generator)
            indices[:] = torch.tensor(indices)[idx].tolist()

    base_batch_sampler_iters = [iter(batch_sampler) for batch_sampler in self.base_batch_samplers]
    bucket_remaining_elements = self.bucket_sizes.clone()
    total_remaining_elements = self.num_samples

    while total_remaining_elements > 0:
        if self.shuffle:
            bucket_idx = torch.multinomial(
                bucket_remaining_elements / total_remaining_elements, 1, generator=self.generator
            )
        else:
            bucket_idx = torch.argmax((bucket_remaining_elements > 0).to(int))  # type: ignore

        try:
            batch = next(base_batch_sampler_iters[bucket_idx])
            bucket_remaining_elements[bucket_idx] -= len(batch)
            total_remaining_elements -= len(batch)
            yield batch
        except StopIteration:
            bucket_remaining_elements[bucket_idx] = 0
            total_remaining_elements = torch.sum(bucket_remaining_elements)
            continue

__len__()

Get the number of batches.

Can only be called if the base_batch_sampler_class has len() implemented

Returns:

Name	Type	Description
int	int	Number of batches

Source code in bionemo/size_aware_batching/sampler.py

def __len__(self) -> int:
    """Get the number of batches.

    Can only be called if the `base_batch_sampler_class` has __len__() implemented

    Returns:
        int: Number of batches
    """
    num_batches = sum(len(sampler) for sampler in self.base_batch_samplers)  # type: ignore
    return num_batches

SizeAwareBatchSampler

Bases: Sampler[List[int]]

Varriying-size batching data sampler class that ensures batch size doesn't exceed maximum.

A sampler that batches elements of varying sizes while ensuring that the total size of each batch does not exceed a specified maximum.

This is useful when dealing with datasets where each element has a different size, such as graphs or sequences of varying lengths. The sampler uses a provided sizeof function to determine the size of each element in the dataset and ensures that the total size of each batch does not exceed the specified max_total_size.

---

Examples:

>>> import torch
>>> from bionemo.size_aware_batching.sampler import SizeAwareBatchSampler


>>> # Define a sample dataset with torch.tensor
>>> dataset = [torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor([5, 6]),
...            torch.tensor([7, 8]), torch.tensor([9, 10])]


>>> # Define a function that returns the size of each element in the dataset.
>>> def sizeof(index):
...     return dataset[index].numel()


>>> # Create a SizeAwareBatchSampler with a maximum total batch size of 10.
>>> batch_sampler = SizeAwareBatchSampler(
...     sampler=torch.utils.data.SequentialSampler(dataset),
...     sizeof=sizeof,
...     max_total_size=4
... )


>>> # Iterate over batches of indices that do not exceed the maximum total size.
>>> print(list(batch_sampler))
    [[0, 1], [2, 3], [4]]

Source code in bionemo/size_aware_batching/sampler.py

class SizeAwareBatchSampler(Sampler[List[int]]):
    """Varriying-size batching data sampler class that ensures batch size doesn't exceed maximum.

    A sampler that batches elements of varying sizes while ensuring
    that the total size of each batch does not exceed a specified maximum.

    This is useful when dealing with datasets where each element has a
    different size, such as graphs or sequences of varying lengths.
    The sampler uses a provided `sizeof` function to determine the size
    of each element in the dataset and ensures that the total size of
    each batch does not exceed the specified `max_total_size`.

    ---------

    Examples:
    ```python
    >>> import torch
    >>> from bionemo.size_aware_batching.sampler import SizeAwareBatchSampler


    >>> # Define a sample dataset with torch.tensor
    >>> dataset = [torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor([5, 6]),
    ...            torch.tensor([7, 8]), torch.tensor([9, 10])]


    >>> # Define a function that returns the size of each element in the dataset.
    >>> def sizeof(index):
    ...     return dataset[index].numel()


    >>> # Create a SizeAwareBatchSampler with a maximum total batch size of 10.
    >>> batch_sampler = SizeAwareBatchSampler(
    ...     sampler=torch.utils.data.SequentialSampler(dataset),
    ...     sizeof=sizeof,
    ...     max_total_size=4
    ... )


    >>> # Iterate over batches of indices that do not exceed the maximum total size.
    >>> print(list(batch_sampler))
        [[0, 1], [2, 3], [4]]
    ```
    """

    def __init__(
        self,
        sampler: Union[Sampler[List[int]], Iterable[int]],
        sizeof: Callable[[int], Real],
        max_total_size: Real,
        info_logger: Optional[Callable[[str], None]] = None,
        warn_logger: Optional[Callable[[str], None]] = None,
    ) -> None:
        """Initializes the SizeAwareBatchSampler.

        Args:
            sampler: The underlying sampler.
            sizeof: A function that returns the size at each index. E.g., this can used to
                determine how much memory an element consumes. Its return type must be
                comparable with `max_total_size` and it must be addable (operator `+`).
            max_total_size: The maximum total size of a mini-batch. The semantics of "size"
                is defined by the `sizeof` argument. The type of this value must be comparable
                with the return type of sizeof, i.e., the operator `<` and `==` must be meaningful.
            info_logger: A function to log info. Defaults to None.
            warn_logger: A function to log warnings. Defaults None.

        Raises:
            TypeError: If sampler is not an instance of Sampler or Iterable, or if sizeof is not a callable, dictionary, or sequence container.
            ValueError: If max_total_size is not a positive number.

        """
        if not (isinstance(sampler, Sampler) or (isinstance(sampler, Iterable) and not isinstance(sampler, str))):
            raise TypeError("sampler should be an instance of torch.utils.data.Sampler or Iterable")

        if not isinstance(max_total_size, Real):
            raise ValueError(f"max_total_size should be int or float but got {type(max_total_size)}")

        self._info_logger = info_logger
        self._warn_logger = warn_logger

        self._is_sizeof_callable = callable(sizeof)

        if not self._is_sizeof_callable:
            raise TypeError("sizeof must be a callable")

        self._sampler = sampler
        self._sizeof = sizeof
        self._max_total_size = max_total_size

    def __iter__(self) -> Iterator[List[int]]:
        """Iterate over batches of indices.

        This function yields batches of indices that do not exceed the maximum total size.

        Yields:
            A batch of indices that do not exceed the maximum total size.
        """
        return size_aware_batching(
            self._sampler,
            self._sizeof,
            self._max_total_size,
            collate_fn=None,
            info_logger=self._info_logger,
            warn_logger=self._warn_logger,
        )

__init__(sampler, sizeof, max_total_size, info_logger=None, warn_logger=None)

Initializes the SizeAwareBatchSampler.

Parameters:

Name	Type	Description	Default
sampler	Union[Sampler[List[int]], Iterable[int]]	The underlying sampler.	required
sizeof	Callable[[int], Real]	A function that returns the size at each index. E.g., this can used to determine how much memory an element consumes. Its return type must be comparable with max_total_size and it must be addable (operator +).	required
max_total_size	Real	The maximum total size of a mini-batch. The semantics of "size" is defined by the sizeof argument. The type of this value must be comparable with the return type of sizeof, i.e., the operator < and == must be meaningful.	required
info_logger	Optional[Callable[[str], None]]	A function to log info. Defaults to None.	None
warn_logger	Optional[Callable[[str], None]]	A function to log warnings. Defaults None.	None

Raises:

Type	Description
TypeError	If sampler is not an instance of Sampler or Iterable, or if sizeof is not a callable, dictionary, or sequence container.
ValueError	If max_total_size is not a positive number.

Source code in bionemo/size_aware_batching/sampler.py

def __init__(
    self,
    sampler: Union[Sampler[List[int]], Iterable[int]],
    sizeof: Callable[[int], Real],
    max_total_size: Real,
    info_logger: Optional[Callable[[str], None]] = None,
    warn_logger: Optional[Callable[[str], None]] = None,
) -> None:
    """Initializes the SizeAwareBatchSampler.

    Args:
        sampler: The underlying sampler.
        sizeof: A function that returns the size at each index. E.g., this can used to
            determine how much memory an element consumes. Its return type must be
            comparable with `max_total_size` and it must be addable (operator `+`).
        max_total_size: The maximum total size of a mini-batch. The semantics of "size"
            is defined by the `sizeof` argument. The type of this value must be comparable
            with the return type of sizeof, i.e., the operator `<` and `==` must be meaningful.
        info_logger: A function to log info. Defaults to None.
        warn_logger: A function to log warnings. Defaults None.

    Raises:
        TypeError: If sampler is not an instance of Sampler or Iterable, or if sizeof is not a callable, dictionary, or sequence container.
        ValueError: If max_total_size is not a positive number.

    """
    if not (isinstance(sampler, Sampler) or (isinstance(sampler, Iterable) and not isinstance(sampler, str))):
        raise TypeError("sampler should be an instance of torch.utils.data.Sampler or Iterable")

    if not isinstance(max_total_size, Real):
        raise ValueError(f"max_total_size should be int or float but got {type(max_total_size)}")

    self._info_logger = info_logger
    self._warn_logger = warn_logger

    self._is_sizeof_callable = callable(sizeof)

    if not self._is_sizeof_callable:
        raise TypeError("sizeof must be a callable")

    self._sampler = sampler
    self._sizeof = sizeof
    self._max_total_size = max_total_size

__iter__()

Iterate over batches of indices.

This function yields batches of indices that do not exceed the maximum total size.

Yields:

Type	Description
List[int]	A batch of indices that do not exceed the maximum total size.

Source code in bionemo/size_aware_batching/sampler.py

def __iter__(self) -> Iterator[List[int]]:
    """Iterate over batches of indices.

    This function yields batches of indices that do not exceed the maximum total size.

    Yields:
        A batch of indices that do not exceed the maximum total size.
    """
    return size_aware_batching(
        self._sampler,
        self._sizeof,
        self._max_total_size,
        collate_fn=None,
        info_logger=self._info_logger,
        warn_logger=self._warn_logger,
    )

size_aware_batching(dataset, sizeof, max_total_size, collate_fn=None, info_logger=None, warn_logger=None)

Creates a batching iterator where each batch size varries (within a max limit) according to memory consumption.

A generator that batches elements from an iterable while ensuring that the total size of each batch does not exceed a specified maximum. Here the size can be a measurement of memory consumption of the elements in the batch. This can be useful for both indexible data or non-indexible but iterable data.

Parameters:

Name	Type	Description	Default
dataset	Iterable[Data]	The input iterable.	required
sizeof	Callable[[Data], Real]	A function or mapping that returns the "size" of each element in dataset. E.g., this can used to determine how much memory an element consumes. Its return type must be comparable with max_total_size and it must be addable (operator +).	required
max_total_size	Real	The maximum total "size" of each batch. The semantics of "size" is defined by the sizeof argument. The type of this value must be comparable with the return type of sizeof, i.e., the operator < and == must be meaningful.	required
collate_fn	Optional[Callable[[Iterable[Data]], BatchCollated]]	An optional function to collate batches. Defaults to None, in which case each batch is a list of elements from the input dataset	None
info_logger	Optional[Callable[[str], None]]	A function to log info. Defaults to None.	None
warn_logger	Optional[Callable[[str], None]]	A function to log warnings. Defaults to None.	None

Yields:

Type	Description
Union[List[Data], BatchCollated]	A generator that yields batches from dataset.

---

Assumptions 1. Linear complexity. This function consumes the given Iterable of data (dataset) once, by going over the data item one by one to build a batch and yield it as soon as the addition of the next data item to the batch would exceed max_total_size or if the batch is the last one (end of iteration) 2. Additive size measurement. For the general usage case of building mini-batches with a threshold of the batch's memory consumption, it assumes that the size of the batch is the sum of all elements in the batch (additive property). 3. Comparable type of max_total_size and sizeof's return. sizeof's return values must be compared with max_total_size to threshold the size of batches

---

Caveat 1: The generated batch sizes may have large variance - how to workaround: filter the output of this generator using a batch size threshold 2: The number of batches may vary a lot across different epochs. - how to workaround: increase the number of steps that compose an epoch, e.g., in the Lightning training/validation loop, which effectively increases the input dataset size per epoch

---

Example:

>>> import torch
>>> from torch.utils.data import default_collate
>>> from bionemo.size_aware_batching.sampler import size_aware_batching

>>> # Define a sample dataset with torch.tensor
>>> dataset = [torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor([5, 6]),
...            torch.tensor([7, 8]), torch.tensor([9, 10])]

>>> # Define a sizeof function that returns the size of each tensor
>>> def sizeof(x):
...     return x.numel()

>>> # Create a generator with max_total_size=4 and default_collate_fn
>>> gen = size_aware_batching(dataset, sizeof, 4, collate_fn=default_collate)
>>> batches = list(gen)
>>> print(batches)
    [tensor([[1, 2], [3, 4]]), tensor([[5, 6], [7, 8]]), tensor([[9, 10]])]

Source code in bionemo/size_aware_batching/sampler.py

def size_aware_batching(
    dataset: Iterable[Data],
    sizeof: Callable[[Data], Real],
    max_total_size: Real,
    collate_fn: Optional[Callable[[Iterable[Data]], BatchCollated]] = None,
    info_logger: Optional[Callable[[str], None]] = None,
    warn_logger: Optional[Callable[[str], None]] = None,
) -> Iterator[Union[List[Data], BatchCollated]]:
    """Creates a batching iterator where each batch size varries (within a max limit) according to memory consumption.

    A generator that batches elements from an iterable while ensuring that the
    total size of each batch does not exceed a specified maximum. Here the size
    can be a measurement of memory consumption of the elements in the batch.
    This can be useful for both indexible data or non-indexible but iterable data.

    Args:
        dataset: The input iterable.
        sizeof: A function or mapping that returns the "size" of each element in `dataset`.
            E.g., this can used to determine how much memory an element consumes. Its return
            type must be comparable with `max_total_size` and it must be addable (operator `+`).
        max_total_size: The maximum total "size" of each batch. The semantics of "size"
            is defined by the `sizeof` argument. The type of this value must be comparable
            with the return type of sizeof, i.e., the operator `<` and `==` must be meaningful.
        collate_fn: An optional function to collate batches. Defaults to None, in which case
            each batch is a list of elements from the input dataset
        info_logger: A function to log info. Defaults to None.
        warn_logger: A function to log warnings. Defaults to None.

    Yields:
        A generator that yields batches from `dataset`.

    -----------
    Assumptions
    1. Linear complexity. This function consumes the given Iterable of data (`dataset`) once,
       by going over the data item one by one to build a batch and yield it as soon as the
       addition of the next data item to the batch would exceed `max_total_size` or if the
       batch is the last one (end of iteration)
    2. Additive size measurement. For the general usage case of building mini-batches with
       a threshold of the batch's memory consumption, it assumes that the size of the batch is
       the sum of all elements in the batch (additive property).
    3. Comparable type of `max_total_size` and `sizeof`'s return. `sizeof`'s return values
       must be compared with `max_total_size` to threshold the size of batches


    ------
    Caveat
    1: The generated batch sizes may have large variance
       - how to workaround: filter the output of this generator using a batch size threshold
    2: The number of batches may vary a lot across different epochs.
       - how to workaround: increase the number of steps that compose an epoch,
         e.g., in the Lightning training/validation loop, which effectively increases the input
         dataset size per epoch


    -------

    Example:
    ```python
    >>> import torch
    >>> from torch.utils.data import default_collate
    >>> from bionemo.size_aware_batching.sampler import size_aware_batching

    >>> # Define a sample dataset with torch.tensor
    >>> dataset = [torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor([5, 6]),
    ...            torch.tensor([7, 8]), torch.tensor([9, 10])]

    >>> # Define a sizeof function that returns the size of each tensor
    >>> def sizeof(x):
    ...     return x.numel()

    >>> # Create a generator with max_total_size=4 and default_collate_fn
    >>> gen = size_aware_batching(dataset, sizeof, 4, collate_fn=default_collate)
    >>> batches = list(gen)
    >>> print(batches)
        [tensor([[1, 2], [3, 4]]), tensor([[5, 6], [7, 8]]), tensor([[9, 10]])]
    ```

    """
    is_sizeof_callable = callable(sizeof)
    has_collate_fn = collate_fn is not None and callable(collate_fn)

    if not is_sizeof_callable:
        raise TypeError("sizeof must be a callable")

    batch_total_size = 0
    batch = []
    n_samples = 0
    n_samples_batched = 0
    n_batches = 0
    for data in dataset:
        n_samples += 1
        try:
            new_size = sizeof(data)
        except Exception as e:
            raise RuntimeError(f"sizeof raises error at data={data}: {e}") from e
        if new_size > max_total_size:
            if warn_logger is not None:
                warn_logger(
                    f"Size of element {data} exceeds max_total_size" f" ({new_size} > {max_total_size}), skipping"
                )
            continue
        if new_size + batch_total_size > max_total_size:
            n_batches += 1
            if has_collate_fn:
                yield collate_fn(batch)
            else:
                yield batch
            batch_total_size = 0
            batch = []
        batch.append(data)
        n_samples_batched += 1
        batch_total_size += new_size

    # return the remaining batch if there is
    if len(batch) > 0:
        n_batches += 1
        if has_collate_fn:
            yield collate_fn(batch)
        else:
            yield batch

    if warn_logger is not None and n_samples_batched < n_samples:
        warn_logger(
            f"{n_samples_batched} samples were batched from {n_samples} "
            f"of the input data. Missing samples are due to exceeding max_total_size={max_total_size})"
        )

    if info_logger is not None:
        info_logger(
            f"Batched {n_samples_batched} samples into {n_batches} batches. "
            f"If this doesn't match the your expectation, consider adjusting "
            f"max_total_size or the sizeof functor"
        )