Method make_loader

myriad/data_2d.py:435–496

(
        self,
        shards: str,
        batch_size: int,
        num_workers: int,
        shuffle: int = 0,
    )

        return valid

    def make_loader(
        self,
        shards: str,
        batch_size: int,
        num_workers: int,
        shuffle: int = 0,
    ):
        # resolve base directories
        base_dirs = self.tar_base
        if isinstance(base_dirs, (str, Path)):
            base_dirs = [base_dirs]
        elif isinstance(base_dirs, ListConfig):
            base_dirs = OmegaConf.to_object(base_dirs)
        elif isinstance(base_dirs, list):
            base_dirs = base_dirs
        else:
            raise NotImplementedError(f'Tar base is of type {type(base_dirs)} which is not supported as of now')

        base_dirs = [Path(b).expanduser().resolve() for b in base_dirs]

        # search for shards
        if shards is None:
            shard_urls = [
                str(p) for base in base_dirs for p in base.rglob("*.tar")
            ]
        else:
            if isinstance(shards, ListConfig):
                shards = OmegaConf.to_object(shards)
            if isinstance(shards, (list, tuple)):
                patterns = shards
            else:                             # a single string
                patterns = [shards]

            shard_urls = []
            for base in base_dirs:
                for pat in patterns:
                    full_pat = str(base / pat)
                    matches = glob.glob(full_pat)
                    shard_urls.extend(matches)
        
        if len(shard_urls) == 0:
            raise FileNotFoundError("No shards matched patterns")

        shard_urls = list(set(shard_urls))  # deduplicate
        shard_urls.sort()   # sort

        # data pipeline
        dataset = wds.DataPipeline(
            wds.SimpleShardList(shard_urls),
            wds.detshuffle() if self.shuffle else lambda x: x,
            wds.split_by_node,
            wds.split_by_worker,
            partial(wds.tarfile_samples, handler=wds.warn_and_continue),
            *([wds.shuffle(shuffle)] if shuffle != 0 and self.shuffle else []),
            wds.map(self._decode),
            wds.select(self._filter_valid),
            wds.batched(batch_size, partial=False, collation_fn=dict_collation_fn),
        )