Faster BatchSampler with big batch size
Labels: module: dataloader, triaged

### 🚀 The feature, motivation and pitch

In the issue #76950, didicout made some improvements on `BatchSampler` in `torch.utils.data.sampler`, which greatly boost the speed when `drop_last=True`. I was interested in the code (as follows) of ejguan which tried to speed up when `drop_last=False` but failed. I made some improvements to make it effective.

```python3
def special_next(it):
    try:
        return next(it)
    except StopIteration:
        return PlaceHolderObj()

sampler_iter = iter(self.sampler)
if drop_last:
    ...
else:
    while True:
        batch = [special_next(sampler_iter) for _ in range(self.batch_size)]
        if not isinstance(batch[-1], PlaceHolderObj):
            yield batch
        else:
            idx = 0
            for d in batch:
                if isinstance(d, PlaceHolderObj):
                    break
                idx += 1
            if idx > 0:
                yield batch[0: idx]
            break
```
@ejguan 

### Alternatives

Here are my improved code for case `drop_last=False`

```python3
def __iter__(self) -> Iterator[List[int]]:
    # Implemented based on the benchmarking in https://github.com/pytorch/pytorch/pull/76951
    if self.drop_last:
        sampler_iter = iter(self.sampler)
        while True:
            try:
                batch = [next(sampler_iter) for _ in range(self.batch_size)]
                yield batch
            except StopIteration:
                break
    else:
        batch = [PlaceHolderObj()] * self.batch_size
        sampler_iter = iter(self.sampler)
        try:
            while True:
                for i in range(self.batch_size):
                    batch[i] = next(sampler_iter)
                yield batch
                batch = [PlaceHolderObj()] * self.batch_size
        except StopIteration:
            idx = 0
            for value in batch:
                if isinstance(value, PlaceHolderObj):
                    break
                else:
                    idx += 1
            if idx > 0:
                yield batch[:idx]
```

### Additional context

My test env and setting:
```
os: Ubuntu 22.04
cpu: Intel® Core™ i7 13700k
python: 3.8.16 and 3.10.0
pytorch: 2.0.1
```
Test result of python 3.8.16:
```
  batch_size  drop_last      current(avg)    current(std)    new(avg)    new(std)  speedup(new vs current)
------------  -----------  --------------  --------------  ----------  ----------  -------------------------
           4  False              0.004253       0.0001049    0.005348   6.951e-05  -20.46%
           8  False              0.003916       8.08e-05     0.003916   6.295e-05  0.02%
          16  False              0.003652       6.122e-05    0.003019   3.849e-05  20.95%
          32  False              0.00345        4.221e-05    0.002511   4.816e-05  37.43%
          64  False              0.003246       4.454e-05    0.002201   3.64e-05   47.52%
         128  False              0.003213       5.322e-05    0.002088   3.325e-05  53.89%
         640  False              0.003319       4.307e-05    0.002169   2.648e-05  53.01%
        6400  False              0.00341        4.636e-05    0.002392   6.074e-05  42.53%
```

Test result of python 3.10.0:
```
  batch_size  drop_last      current(avg)    current(std)    new(avg)    new(std)  speedup(new vs current)
------------  -----------  --------------  --------------  ----------  ----------  -------------------------
           4  False              0.004014       0.0001283    0.004929   5.085e-05  -18.57%
           8  False              0.003872       0.0001085    0.003761   5.127e-05  2.95%
          16  False              0.003491       6.523e-05    0.002881   3.742e-05  21.20%
          32  False              0.003201       6.749e-05    0.002447   4.072e-05  30.78%
          64  False              0.00333        8.508e-05    0.002371   5.883e-05  40.45%
         128  False              0.003085       8.315e-05    0.002196   3.641e-05  40.48%
         640  False              0.003074       6.132e-05    0.002207   3.811e-05  39.26%
        6400  False              0.003101       5.371e-05    0.002343   2.617e-05  32.34%
```

The following testing code was adapted from didicout:
```python3
# coding: utf-8
# Copyright (c) Antfin, Inc. All rights reserved.

from typing import List, Union, Iterable, Iterator

from torch.utils.data import SequentialSampler, Sampler



class OldBatchSampler(Sampler[List[int]]):
    r"""Wraps another sampler to yield a mini-batch of indices.

    Args:
        sampler (Sampler or Iterable): Base sampler. Can be any iterable object
        batch_size (int): Size of mini-batch.
        drop_last (bool): If ``True``, the sampler will drop the last batch if
            its size would be less than ``batch_size``

    Example:
        >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False))
        [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
        >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True))
        [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
    """

    def __init__(self, sampler: Union[Sampler[int], Iterable[int]], batch_size: int, drop_last: bool) -> None:
        # Since collections.abc.Iterable does not check for `__getitem__`, which
        # is one way for an object to be an iterable, we don't do an `isinstance`
        # check here.
        if not isinstance(batch_size, int) or isinstance(batch_size, bool) or \
                batch_size <= 0:
            raise ValueError("batch_size should be a positive integer value, "
                             "but got batch_size={}".format(batch_size))
        if not isinstance(drop_last, bool):
            raise ValueError("drop_last should be a boolean value, but got "
                             "drop_last={}".format(drop_last))
        self.sampler = sampler
        self.batch_size = batch_size
        self.drop_last = drop_last

    def __iter__(self) -> Iterator[List[int]]:
        # Implemented based on the benchmarking in https://github.com/pytorch/pytorch/pull/76951

        if self.drop_last:
            sampler_iter = iter(self.sampler)
            while True:
                try:
                    batch = [next(sampler_iter) for _ in range(self.batch_size)]
                    yield batch
                except StopIteration:
                    break
        else:

            batch = [0] * self.batch_size
            idx_in_batch = 0
            for idx in self.sampler:
                batch[idx_in_batch] = idx
                idx_in_batch += 1
                if idx_in_batch == self.batch_size:
                    yield batch
                    idx_in_batch = 0
                    batch = [0] * self.batch_size
            if idx_in_batch > 0:
                yield batch[:idx_in_batch]

    def __len__(self) -> int:
        # Can only be called if self.sampler has __len__ implemented
        # We cannot enforce this condition, so we turn off typechecking for the
        # implementation below.
        # Somewhat related: see NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ]
        if self.drop_last:
            return len(self.sampler) // self.batch_size  # type: ignore[arg-type]
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size  # type: ignore[arg-type]


class PlaceHolderObj:
    pass


class NewBatchSampler(Sampler[List[int]]):
    def __init__(self, sampler: Union[Sampler[int], Iterable[int]], batch_size: int, drop_last: bool) -> None:
        # Since collections.abc.Iterable does not check for `__getitem__`, which
        # is one way for an object to be an iterable, we don't do an `isinstance`
        # check here.
        if not isinstance(batch_size, int) or isinstance(batch_size, bool) or \
                batch_size <= 0:
            raise ValueError("batch_size should be a positive integer value, "
                             "but got batch_size={}".format(batch_size))
        if not isinstance(drop_last, bool):
            raise ValueError("drop_last should be a boolean value, but got "
                             "drop_last={}".format(drop_last))
        self.sampler = sampler
        self.batch_size = batch_size
        self.drop_last = drop_last

    def __iter__(self) -> Iterator[List[int]]:
        # Implemented based on the benchmarking in https://github.com/pytorch/pytorch/pull/76951
        if self.drop_last:
            sampler_iter = iter(self.sampler)
            while True:
                try:
                    batch = [next(sampler_iter) for _ in range(self.batch_size)]
                    yield batch
                except StopIteration:
                    break
        else:
            batch = [PlaceHolderObj()] * self.batch_size
            sampler_iter = iter(self.sampler)
            try:
                while True:
                    for i in range(self.batch_size):
                        batch[i] = next(sampler_iter)
                    yield batch
                    batch = [PlaceHolderObj()] * self.batch_size
            except StopIteration:
                idx = 0
                for value in batch:
                    if isinstance(value, PlaceHolderObj):
                        break
                    else:
                        idx += 1
                if idx > 0:
                    yield batch[:idx]

    def __len__(self) -> int:
        # Can only be called if self.sampler has __len__ implemented
        # We cannot enforce this condition, so we turn off typechecking for the
        # implementation below.
        # Somewhat related: see NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ]
        if self.drop_last:
            return len(self.sampler) // self.batch_size  # type: ignore[arg-type]
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size  # type: ignore[arg-type]


def _iter_on_current_sampler(batch_size, drop_last):
    for _ in OldBatchSampler(SequentialSampler(range(DATA_SIZE)), batch_size=batch_size, drop_last=drop_last):
        pass


def _iter_on_new_sampler(batch_size, drop_last):
    for _ in NewBatchSampler(SequentialSampler(range(DATA_SIZE)), batch_size=batch_size, drop_last=drop_last):
        pass


if __name__ == '__main__':
    import numpy as np
    import timeit
    from tabulate import tabulate

    DATA_SIZE = 99999
    AVG_TIMES = 100

    result = []
    DROP_LAST = False
    for BATCH_SIZE in (4, 8, 16, 32, 64, 128, 640, 6400):
        print(f"batch size: {BATCH_SIZE}, drop last: {DROP_LAST}")

        timer = timeit.Timer(lambda: _iter_on_current_sampler(BATCH_SIZE, DROP_LAST))
        cost0 = timer.repeat(AVG_TIMES, 1)
        cost0_avg = float(np.average(cost0))
        cost0_std = float(np.std(cost0))
        print(f"time cost(current):{cost0_avg}. origin data: {cost0}")

        timer = timeit.Timer(lambda: _iter_on_new_sampler(BATCH_SIZE, DROP_LAST))
        cost1 = timer.repeat(AVG_TIMES, 1)
        cost1_avg = float(np.average(cost1))
        cost1_std = float(np.std(cost1))
        print(f"time cost(new):{cost1_avg}. origin data: {cost1}")

        speedup_percent = "%.2f" % ((1 / cost1_avg - 1 / cost0_avg) * cost0_avg * 100) + "%"
        print(f"speedup: {speedup_percent}\n")

        current_row = [BATCH_SIZE, DROP_LAST,
                        "%.3e" % cost0_avg, "%.3e" % cost0_std,
                        "%.3e" % cost1_avg, "%.3e" % cost1_std,
                        speedup_percent]

        result.append(current_row)

    print(tabulate(result, headers=("batch_size", "drop_last", "current(avg)", "current(std)", "new(avg)", "new(std)",
                                    "speedup(new vs current)")))
```

cc @SsnL @VitalyFedyunin @ejguan @NivekT @dzhulgakov