Repository: Project-MONAI/MONAI
File: monai/data/grid_dataset.py (Python)

Code before:
        if indices is None:
            indices = list(range(self.cache_num))
        if self.progress and not has_tqdm:
            warnings.warn("tqdm is not installed, will not show the caching progress bar.")
        with ThreadPool(self.num_workers) as p:
            if self.progress and has_tqdm:
                return list(tqdm(p.imap(self._load_cache_item, indices), total=len(indices), desc="Loading dataset"))
            return list(p.imap(self._load_cache_item, indices))

    def _load_cache_item(self, idx: int):
        """
        Args:
            idx: the index of the input data sequence.
        """
        item = self.data[idx]  # type: ignore
        patch_cache, other_cache = [], []
        for patch, *others in self.patch_iter(item):
            if self.first_random is not None:
                patch = self.patch_transform(patch, end=self.first_random, threading=True)  # type: ignore

            if self.as_contiguous:
                patch = convert_to_contiguous(patch, memory_format=torch.contiguous_format)
            if self.with_coordinates and len(others) > 0:  # patch_iter to yield at least 2 items: patch, coords
                other_cache.append(others[0])
            patch_cache.append(patch)
        self._cache_other.append(other_cache)
        return patch_cache

    def _generate_patches(self, src, **apply_args):
        """
        yield patches optionally post-processed by transform.

        Args:
            src: a iterable of image patches.
            apply_args: other args for `self.patch_transform`.

        """
        for patch, *others in src:
            out_patch = patch
            if self.patch_transform is not None:
                out_patch = self.patch_transform(patch, **apply_args)
            if self.with_coordinates and len(others) > 0:  # patch_iter to yield at least 2 items: patch, coords
                yield out_patch, others[0]
            else:
                yield out_patch

    def __iter__(self):
        if self.cache:
            cache_index = None
            for image in super().__iter__():
                key = self.hash_func(image)

Review comment (performance):
This isn't thread-safe since `_cache_other` is just a `list` object and so there's no locking. There's also no guarantee of ordering if the threads in the `ThreadPool` operate at different timings than what you've tested with. I think you need to be accumulating a `cache_other` list at the same time you're accumulating `patch_cache` and also return it, then in `set_data` have something like `self._cache, self._cache_other = zip(*self._fill_cache(indices))` to produce the two lists that you need. 

Code after:

        pfunc = tqdm if self.progress and has_tqdm else (lambda v, **_: v)
        with ThreadPool(self.num_workers) as p:
            return list(pfunc(p.imap(self._load_cache_item, indices), total=len(indices), desc="Loading dataset"))

    def _load_cache_item(self, idx: int):
        """
        Args:
            idx: the index of the input data sequence.
        """
        item = self.data[idx]  # type: ignore
        patch_cache, other_cache = [], []
        for patch, *others in self.patch_iter(item):
            if self.first_random is not None:
                patch = self.patch_transform(patch, end=self.first_random, threading=True)  # type: ignore

            if self.as_contiguous:
                patch = convert_to_contiguous(patch, memory_format=torch.contiguous_format)
            if self.with_coordinates and len(others) > 0:  # patch_iter to yield at least 2 items: patch, coords
                other_cache.append(others[0])
            patch_cache.append(patch)
        return patch_cache, other_cache

    def _generate_patches(self, src, **apply_args):
        """
        yield patches optionally post-processed by transform.

        Args:
            src: a iterable of image patches.
            apply_args: other args for `self.patch_transform`.

        """
        for patch, *others in src:
            out_patch = patch
            if self.patch_transform is not None:
                out_patch = self.patch_transform(patch, **apply_args)
            if self.with_coordinates and len(others) > 0:  # patch_iter to yield at least 2 items: patch, coords
                yield out_patch, others[0]
            else:
                yield out_patch

    def __iter__(self):
        if self.cache:
            cache_index = None
            for image in super().__iter__():
                key = self.hash_func(image)
                if key in self._hash_keys:
                    # if existing in cache, try to get the index in cache
                    cache_index = self._hash_keys.index(key)
                if cache_index is None:
                    # no cache for this index, execute all the transforms directly