(self)
| 89 | self.num_bucket_build_workers = num_bucket_build_workers |
| 90 | |
| 91 | def __iter__(self) -> Iterator[List[int]]: |
| 92 | if self._get_num_batch_cached_bucket_sample_dict is not None: |
| 93 | bucket_sample_dict = self._get_num_batch_cached_bucket_sample_dict |
| 94 | self._get_num_batch_cached_bucket_sample_dict = None |
| 95 | else: |
| 96 | bucket_sample_dict = self.group_by_bucket() |
| 97 | if self.verbose: |
| 98 | self._print_bucket_info(bucket_sample_dict) |
| 99 | |
| 100 | g = torch.Generator() |
| 101 | g.manual_seed(self.seed + self.epoch) |
| 102 | bucket_micro_batch_count = OrderedDict() |
| 103 | bucket_last_consumed = OrderedDict() |
| 104 | |
| 105 | # process the samples |
| 106 | for bucket_id, data_list in bucket_sample_dict.items(): |
| 107 | # handle droplast |
| 108 | bs_per_gpu = self.bucket.get_batch_size(bucket_id) |
| 109 | remainder = len(data_list) % bs_per_gpu |
| 110 | |
| 111 | if remainder > 0: |
| 112 | if not self.drop_last: |
| 113 | # if there is remainder, we pad to make it divisible |
| 114 | data_list += data_list[:bs_per_gpu - remainder] |
| 115 | else: |
| 116 | # we just drop the remainder to make it divisible |
| 117 | data_list = data_list[:-remainder] |
| 118 | bucket_sample_dict[bucket_id] = data_list |
| 119 | |
| 120 | # handle shuffle |
| 121 | if self.shuffle: |
| 122 | data_indices = torch.randperm( |
| 123 | len(data_list), generator=g).tolist() |
| 124 | data_list = [data_list[i] for i in data_indices] |
| 125 | bucket_sample_dict[bucket_id] = data_list |
| 126 | |
| 127 | # compute how many micro-batches each bucket has |
| 128 | num_micro_batches = len(data_list) // bs_per_gpu |
| 129 | bucket_micro_batch_count[bucket_id] = num_micro_batches |
| 130 | |
| 131 | # compute the bucket access order |
| 132 | # each bucket may have more than one batch of data |
| 133 | # thus bucket_id may appear more than 1 time |
| 134 | bucket_id_access_order = [] |
| 135 | for bucket_id, num_micro_batch in bucket_micro_batch_count.items(): |
| 136 | bucket_id_access_order.extend([bucket_id] * num_micro_batch) |
| 137 | |
| 138 | # randomize the access order |
| 139 | if self.shuffle: |
| 140 | bucket_id_access_order_indices = torch.randperm( |
| 141 | len(bucket_id_access_order), generator=g).tolist() |
| 142 | bucket_id_access_order = [ |
| 143 | bucket_id_access_order[i] |
| 144 | for i in bucket_id_access_order_indices |
| 145 | ] |
| 146 | |
| 147 | # make the number of bucket accesses divisible by dp size |
| 148 | remainder = len(bucket_id_access_order) % self.num_replicas |
nothing calls this directly
no test coverage detected