Method __iter__

data/dataset_base_periter.py:308–383

(self)

        return data

    def __iter__(self):

        num_groups = len(self.dataset_iters)
        ################## group weight ############
        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))]

        ##############################
        while True:
            sequence_status = self.set_sequence_status()
            batch_data_indexes = []
            added_at_least_one = False
    
            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
            )      

            n = step_rng.random() # Use the step_rng for reproducibility
            group_index = 0
            for i, cumprob in enumerate(group_cumprobs):
                if n < cumprob:
                    group_index = i
                    break
            group_iter, group_name, group_dataset = self.dataset_iters[group_index]

            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
                # if a sample is too long, skip it
                num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan']) 

                if num_tokens == 0 or num_tokens > self.max_num_tokens_per_sample:
                    if num_tokens == 0:
                        print(f"skip a sample with length 0")
                    else:
                        print(f"skip a sample with length {num_tokens} (exceeds max_per_sample)")
                    continue