Method cache_latents

utils/dataset.py:233–298

(self, map_fn, regenerate_cache=False, trust_cache=False, caching_batch_size=1)

            raise ValueError(f'num_repeats must be >0, was {self.num_repeats}')

    def cache_latents(self, map_fn, regenerate_cache=False, trust_cache=False, caching_batch_size=1):
        print(f'caching latents: {self.size_bucket}')
        self.latent_dataset = _map_and_cache(
            self.metadata_dataset,
            map_fn,
            self.cache_dir,
            cache_file_prefix='latents_',
            regenerate_cache=regenerate_cache,
            caching_batch_size=caching_batch_size,
        )
        assert len(self.latent_dataset) == len(self.metadata_dataset), (len(self.latent_dataset), len(self.metadata_dataset))

        iteration_order_cache_dir = self.cache_dir / 'iteration_order'

        if regenerate_cache or not iteration_order_cache_dir.exists() or not trust_cache:
            print('Building iteration order')
            image_spec_to_latents_idx = {
                tuple(image_spec): i
                for i, image_spec in enumerate(self.metadata_dataset['image_spec'])
            }

            equal_num_captions = True
            num_captions = None
            for example in self.metadata_dataset.select_columns(['caption']):
                n = len(example['caption'])
                if num_captions is not None and n != num_captions:
                    equal_num_captions = False
                    break
                num_captions = n

            if equal_num_captions:
                # If all images have the same number of captions, set things up so we read (mostly) sequentially off disk. The metadata was already shuffled in the beginning.
                iteration_order_by_caption_num = [[] for _ in range(num_captions)]
                seed = 0
                for example in self.metadata_dataset.select_columns(['image_spec', 'caption']):
                    image_spec = example['image_spec']
                    captions = example['caption']
                    shuffle_with_seed(captions, seed)
                    seed += 1
                    latents_idx = image_spec_to_latents_idx[tuple(image_spec)]
                    for i, caption in enumerate(captions):
                        iteration_order_by_caption_num[i].append((image_spec, latents_idx, caption, i))
                iteration_order_list = []
                for l in iteration_order_by_caption_num:
                    iteration_order_list.extend(l)
            else:
                iteration_order_list = []
                for example in self.metadata_dataset.select_columns(['image_spec', 'caption']):
                    image_spec = example['image_spec']
                    captions = example['caption']
                    latents_idx = image_spec_to_latents_idx[tuple(image_spec)]
                    for i, caption in enumerate(captions):
                        iteration_order_list.append((image_spec, latents_idx, caption, i))
                shuffle_with_seed(iteration_order_list, 42)

            iteration_order_dict = defaultdict(list)
            for image_spec, latents_idx, caption, caption_number in iteration_order_list:
                iteration_order_dict['image_spec'].append(image_spec)