Sampler

BucketBatchSampler

基类：Sampler[List[int]]

一个批采样器，用于创建批次，其中批次中元素的大小来自每个预定义的桶范围。

数据集的元素首先根据桶范围和元素的大小分组到每个桶中。然后，为每个桶使用一个基本批采样器来创建小批量。

桶范围由 bucket_boundaries 指定，它将首先在内部排序，并用于创建 len(bucket_boundaries) - 1 个左闭右开区间。例如，如果 bucket_boundaries 张量为 [10, 5, 0, 16]，它将被排序为 [0, 5, 10, 16]，并将创建 3 个桶，范围为：[0, 5)、[5, 10)、[10, 16)。

基本批采样器将通过传递每个桶中的元素索引作为数据源，并将 base_batch_sampler_shared_kwargs 和 base_batch_sampler_individual_kwargs 传递给指定为 base_batch_sampler_class 的基本批采样器类的构造函数来创建。例如，base_batch_sampler_shared_kwargs = {'drop_last': True} 和 base_batch_sampler_individual_kwargs = {'batch_size': [8,10,12]} 将用于创建 3 个批采样器，其中 drop_last=True，batch_size 分别为 8、10 和 12，并像 base_batch_sampler_class(bucket_element_indices[0], batch_size=8, drop_last=True) 这样初始化。

在 __iter__ 方法中，如果 shuffle 为 True，则每个桶中的元素索引将被洗牌，并且每次随机选择一个桶来创建小批量。如果 shuffle 为 False，则不会对元素索引进行洗牌，并且桶将按其区间边界的升序选择。

此类用于创建同质数据批次，用于训练或评估，并减少对齐元素形状所需的填充。

修改自 https://github.com/rssrwn/semla-flow/blob/main/semlaflow/data/util.py

---

示例

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

源代码位于 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}"
            )

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

初始化 BucketBatchSampler。

参数

名称	类型	描述	默认值
sizes	Tensor	表示数据集中每个元素大小的实数一维张量。	必需
bucket_boundaries	Tensor	表示桶范围边界的实数一维张量。它将首先被排序，并用于创建 len(bucket_boundaries) - 1 个左闭右开区间作为桶范围。它不应包含任何重复值。	必需
base_batch_sampler_class	Type[S]	基本批采样器类类型，它将用于每个桶，并使用桶元素索引、base_batch_sampler_shared_kwargs 和相应的 base_batch_sampler_individual_kwargs 进行初始化。	必需
base_batch_sampler_shared_kwargs	Optional[Dict[str, Any]]	用于初始化所有桶的所有基本批采样器的共享关键字参数字典。应为 base_batch_sampler_class 提供足够且有效的参数以及 base_batch_sampler_individual_kwargs。默认为 {}。	None
base_batch_sampler_individual_kwargs	Optional[Dict[str, Iterable]]	用于使用相应的键值对初始化每个桶批采样器的关键字参数字典。此字典中每个值的长度必须等于 len(bucket_boundaries) - 1（桶的数量）。应为 base_batch_sampler_class 提供足够且有效的参数以及 base_batch_sampler_shared_kwargs。默认为 {}。	None
shuffle	Optional[bool]	一个布尔值，指示是否洗牌数据集和桶。默认为 True。	True
generator	Optional[Generator]	采样中使用的生成器。默认为 None。	None

引发

类型	描述
ValueError	如果 sizes 不是实数的一维张量。
ValueError	如果 bucket_boundaries 不是实数的一维张量。
ValueError	如果 base_batch_sampler_shared_kwargs 或 base_batch_sampler_individual_kwargs 不是关键字参数字典。
ValueError	如果 base_batch_sampler_individual_kwargs 字典中值的长度必须等于 len(bucket_boundaries) - 1。
RuntimeError	如果桶范围 bucket_boundaries 指定的范围内没有大小合适的元素。

源代码位于 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}"
        )

    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
    )
    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__()

迭代索引批次。

此函数产生来自每个桶范围的大小元素的索引批次。

产生

类型	描述
List[int]	List[int]：来自每个桶范围的大小元素的索引批次。

源代码位于 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__()

获取批次数量。

只有在 base_batch_sampler_class 实现了 len() 时才能调用

返回

名称	类型	描述
int	int	批次数量

源代码位于 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

基类：Sampler[List[int]]

可变大小批处理数据采样器类，确保批次大小不超过最大值。

一种采样器，用于批处理大小不一的元素，同时确保每个批次的总大小不超过指定的最大值。

这在处理每个元素大小不同的数据集时非常有用，例如图形或长度不一的序列。采样器使用提供的 sizeof 函数来确定数据集中每个元素的大小，并确保每个批次的总大小不超过指定的 max_total_size。

---

示例

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

源代码位于 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)

初始化 SizeAwareBatchSampler。

参数

名称	类型	描述	默认值
sampler	Union[Sampler[List[int]], Iterable[int]]	底层采样器。	必需
sizeof	Callable[[int], Real]	一个函数，它返回每个索引处的大小。例如，这可以用于确定一个元素消耗多少内存。其返回类型必须与 max_total_size 可比较，并且必须是可加的（运算符 +）。	必需
max_total_size	Real	小批量的最大总大小。“大小”的语义由 sizeof 参数定义。此值的类型必须与 sizeof 的返回类型可比较，即运算符 < 和 == 必须有意义。	必需
info_logger	Optional[Callable[[str], None]]	用于记录信息的函数。默认为 None。	None
warn_logger	Optional[Callable[[str], None]]	用于记录警告的函数。默认为 None。	None

引发

类型	描述
TypeError	如果采样器不是 Sampler 或 Iterable 的实例，或者如果 sizeof 不是可调用对象、字典或序列容器。
ValueError	如果 max_total_size 不是正数。

源代码位于 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__()

迭代索引批次。

此函数产生不超过最大总大小的索引批次。

产生

类型	描述
List[int]	不超过最大总大小的索引批次。

源代码位于 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)

创建一个批处理迭代器，其中每个批次大小根据内存消耗而变化（在最大限制内）。

一个生成器，用于从可迭代对象中批处理元素，同时确保每个批次的总大小不超过指定的最大值。这里的大小可以是批次中元素内存消耗的度量。这对于可索引数据或不可索引但可迭代的数据都很有用。

参数

名称	类型	描述	默认值
dataset	Iterable[Data]	输入可迭代对象。	必需
sizeof	Callable[[Data], Real]	一个函数或映射，它返回 dataset 中每个元素的“大小”。例如，这可以用于确定一个元素消耗多少内存。其返回类型必须与 max_total_size 可比较，并且必须是可加的（运算符 +）。	必需
max_total_size	Real	每个批次的最大总“大小”。“大小”的语义由 sizeof 参数定义。此值的类型必须与 sizeof 的返回类型可比较，即运算符 < 和 == 必须有意义。	必需
collate_fn	Optional[Callable[[Iterable[Data]], BatchCollated]]	一个可选的函数，用于整理批次。默认为 None，在这种情况下，每个批次都是来自输入数据集的元素列表	None
info_logger	Optional[Callable[[str], None]]	用于记录信息的函数。默认为 None。	None
warn_logger	Optional[Callable[[str], None]]	用于记录警告的函数。默认为 None。	None

产生

类型	描述
Union[List[Data], BatchCollated]	一个从 dataset 产生批次的生成器。

---

假设 1. 线性复杂度。此函数消耗给定的数据 Iterable (dataset) 一次，通过逐个遍历数据项来构建批次，并在将下一个数据项添加到批次中会超过 max_total_size 或如果批次是最后一个批次（迭代结束）时立即产生批次。2. 可加性大小测量。对于构建具有批次内存消耗阈值的小批量的通用用例，它假设批次的大小是批次中所有元素的总和（可加性）。3. max_total_size 和 sizeof 返回的可比较类型。 sizeof 的返回值必须与 max_total_size 进行比较，以阈值化批次的大小

---

注意事项 1：生成的批次大小可能具有很大的差异 - 如何解决：使用批次大小阈值过滤此生成器的输出 2：不同 epoch 之间的批次数量可能会有很大差异。 - 如何解决：增加构成一个 epoch 的步骤数，例如，在 Lightning 训练/验证循环中，这有效地增加了每个 epoch 的输入数据集大小

---

示例

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

源代码位于 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"
        )