goathead777/Zero_Shot_Inference

1# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/bucketsampler.py2import itertools3import math4import random5from random import shuffle6from typing import Iterator7from typing import Optional8from typing import TypeVar9 10import torch11import torch.distributed as dist12from torch.utils.data import Dataset13from torch.utils.data import Sampler14 15__all__ = [16    "DistributedBucketSampler",17]18 19T_co = TypeVar("T_co", covariant=True)20 21 22class DistributedBucketSampler(Sampler[T_co]):23    r"""24    sort the dataset wrt. input length25    divide samples into buckets26    sort within buckets27    divide buckets into batches28    sort batches29    """30 31    def __init__(32        self,33        dataset: Dataset,34        num_replicas: Optional[int] = None,35        rank: Optional[int] = None,36        shuffle: bool = True,37        seed: int = 0,38        drop_last: bool = False,39        batch_size: int = 32,40    ) -> None:41        if num_replicas is None:42            if not dist.is_available():43                raise RuntimeError("Requires distributed package to be available")44            num_replicas = dist.get_world_size()45        if rank is None:46            if not dist.is_available():47                raise RuntimeError("Requires distributed package to be available")48            rank = dist.get_rank()49            torch.cuda.set_device(rank)50        if rank >= num_replicas or rank < 0:51            raise ValueError(52                "Invalid rank {}, rank should be in the interval"53                " [0, {}]".format(rank, num_replicas - 1)54            )55        self.dataset = dataset56        self.num_replicas = num_replicas57        self.rank = rank58        self.epoch = 059        self.drop_last = drop_last60        # If the dataset length is evenly divisible by # of replicas, then there61        # is no need to drop any data, since the dataset will be split equally.62        if (63            self.drop_last and len(self.dataset) % self.num_replicas != 064        ):  # type: ignore[arg-type]65            # Split to nearest available length that is evenly divisible.66            # This is to ensure each rank receives the same amount of data when67            # using this Sampler.68            self.num_samples = math.ceil(69                (len(self.dataset) - self.num_replicas)70                / self.num_replicas  # type: ignore[arg-type]71            )72        else:73            self.num_samples = math.ceil(74                len(self.dataset) / self.num_replicas75            )  # type: ignore[arg-type]76        self.total_size = self.num_samples * self.num_replicas77        self.shuffle = shuffle78        self.seed = seed79        self.batch_size = batch_size80        self.id_with_length = self._get_sample_lengths()81        self.id_buckets = self.make_buckets(bucket_width=2.0)82 83    def _get_sample_lengths(self):84        id_with_lengths = []85        for i in range(len(self.dataset)):86            id_with_lengths.append((i, self.dataset.get_sample_length(i)))87        id_with_lengths.sort(key=lambda x: x[1])88        return id_with_lengths89 90    def make_buckets(self, bucket_width: float = 2.0):91        buckets = []92        cur = []93        max_sec = bucket_width94        for id, sec in self.id_with_length:95            if sec < max_sec:96                cur.append(id)97            else:98                buckets.append(cur)99                cur = [id]100                max_sec += bucket_width101        if len(cur) > 0:102            buckets.append(cur)103        return buckets104 105    def __iter__(self) -> Iterator[T_co]:106        if self.shuffle:107            # deterministically shuffle based on epoch and seed108            g = torch.Generator()109            g.manual_seed(self.seed + self.epoch)110            random.seed(self.epoch + self.seed)111            shuffled_bucket = []112            for buc in self.id_buckets:113                buc_copy = buc.copy()114                shuffle(buc_copy)115                shuffled_bucket.append(buc_copy)116            grouped_batch_size = self.batch_size * self.num_replicas117            shuffled_bucket = list(itertools.chain(*shuffled_bucket))118            n_batch = int(math.ceil(len(shuffled_bucket) / grouped_batch_size))119            batches = [120                shuffled_bucket[b * grouped_batch_size : (b + 1) * grouped_batch_size]121                for b in range(n_batch)122            ]123            shuffle(batches)124            indices = list(itertools.chain(*batches))125        else:126            # type: ignore[arg-type]127            indices = list(range(len(self.dataset)))128 129        if not self.drop_last:130            # add extra samples to make it evenly divisible131            padding_size = self.total_size - len(indices)132            if padding_size <= len(indices):133                indices += indices[:padding_size]134            else:135                indices += (indices * math.ceil(padding_size / len(indices)))[136                    :padding_size137                ]138        else:139            # remove tail of data to make it evenly divisible.140            indices = indices[: self.total_size]141        assert len(indices) == self.total_size142 143        # subsample144        indices = indices[self.rank : self.total_size : self.num_replicas]145        assert len(indices) == self.num_samples146 147        return iter(indices)148 149    def __len__(self) -> int:150        return self.num_samples151 152    def set_epoch(self, epoch: int) -> None:153        r"""154        Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas155        use a different random ordering for each epoch. Otherwise, the next iteration of this156        sampler will yield the same ordering.157 158        Args:159            epoch (int): Epoch number.160        """161        self.epoch = epoch162