Method __iter__

data/dataset_base.py:308–388

(self)

        return data

    def __iter__(self):
        total_weights = sum(self.grouped_weights)
        assert total_weights > 0.0
        group_cumprobs = [sum(self.grouped_weights[:i + 1]) / total_weights 
                        for i in range(len(self.grouped_weights))]
        sequence_status = self.set_sequence_status()
        batch_data_indexes = []

        while True:
            self.step_counter += 1
            step_seed = self.base_and_epoch_seed + self.step_counter
            step_rng = random.Random(step_seed)
            
            random_image_num = int(step_rng.choices(
                self.possible_nums, 
                weights=self.normalized_weights,
                k=1
            )[0])
            random_aspect_ratio = round(
                step_rng.uniform(self.aspect_ratio_range[0], self.aspect_ratio_range[1]), 
                2
            )      

            # Ensure at least one sample from each group
            if sequence_status['curr'] == 0:
                for group_index, group_iter_names in enumerate(self.dataset_iters):
                    group_iter = group_iter_names[0]
                    group_name = group_iter_names[1]
                    group_dataset = group_iter_names[2]  
                    if self.is_mandatory[group_index]: #grouped means recon group.  
                        while True:
                            if group_name == "recon":
                                group_dataset.set_random_image_num(random_image_num)
                                group_dataset.set_random_aspect_ratio(random_aspect_ratio)
                                group_dataset.set_step_rng(step_seed)
                                sample = next(group_iter)
                            else:
                                sample = next(group_iter)
                                
                            if sample is None:
                                continue
                            num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan']) 
                            if num_tokens < self.max_num_tokens_per_sample:
                                sequence_status = self.pack_sequence(sample, sequence_status)
                                batch_data_indexes.append(sample['data_indexes'])
                                break
                            else:
                                print(f"skip a sample with length {num_tokens}")
                                continue

            n = random.random()
            group_index = 0
            for i, cumprob in enumerate(group_cumprobs):
                if n < cumprob:
                    group_index = i
                    break
            sample = next(self.dataset_iters[group_index][0])