| 90 | self.rm_keys = remove_keys if remove_keys is not None else [] |
| 91 | |
| 92 | def make_loader(self, dataset_config, train=True): |
| 93 | image_transforms = [] |
| 94 | lambda_fn = lambda x: x * 2. - 1. # normalize to [-1, 1] |
| 95 | image_transforms.extend([torchvision.transforms.ToTensor(), |
| 96 | torchvision.transforms.Lambda(lambda_fn)]) |
| 97 | if 'image_transforms' in dataset_config: |
| 98 | image_transforms.extend([instantiate_from_config(tt) for tt in dataset_config.image_transforms]) |
| 99 | image_transforms = torchvision.transforms.Compose(image_transforms) |
| 100 | |
| 101 | if 'transforms' in dataset_config: |
| 102 | transforms_config = OmegaConf.to_container(dataset_config.transforms) |
| 103 | else: |
| 104 | transforms_config = dict() |
| 105 | |
| 106 | transform_dict = {dkey: load_partial_from_config(transforms_config[dkey]) |
| 107 | if transforms_config[dkey] != 'identity' else identity |
| 108 | for dkey in transforms_config} |
| 109 | # this is crucial to set correct image key to get the transofrms applied correctly |
| 110 | img_keys = dataset_config.get('image_key', 'image.png') |
| 111 | if isinstance(img_keys, str): |
| 112 | img_keys = [img_keys] |
| 113 | for img_key in img_keys: |
| 114 | transform_dict.update({img_key: image_transforms}) |
| 115 | |
| 116 | if 'dataset_transforms' in dataset_config: |
| 117 | dataset_transforms = instantiate_from_config(dataset_config['dataset_transforms']) |
| 118 | else: |
| 119 | dataset_transforms = None |
| 120 | |
| 121 | if 'postprocess' in dataset_config: |
| 122 | postprocess = instantiate_from_config(dataset_config['postprocess']) |
| 123 | else: |
| 124 | postprocess = None |
| 125 | |
| 126 | shuffle = dataset_config.get('shuffle', 0) |
| 127 | shardshuffle = shuffle > 0 |
| 128 | |
| 129 | nodesplitter = wds.shardlists.split_by_node if self.multinode else wds.shardlists.single_node_only |
| 130 | |
| 131 | if isinstance(dataset_config.shards, str): |
| 132 | tars = os.path.join(self.tar_base, dataset_config.shards) |
| 133 | elif isinstance(dataset_config.shards, list) or isinstance(dataset_config.shards, ListConfig): |
| 134 | # decompose into lists of shards |
| 135 | # Turn train-{000000..000002}.tar into ['train-000000.tar', 'train-000001.tar', 'train-000002.tar'] |
| 136 | tars = [] |
| 137 | for shard in dataset_config.shards: |
| 138 | # Assume that the shard starts from 000000 |
| 139 | if '{' in shard: |
| 140 | start, end = shard.split('..') |
| 141 | start = start.split('{')[-1] |
| 142 | end = end.split('}')[0] |
| 143 | start = int(start) |
| 144 | end = int(end) |
| 145 | tars.extend([shard.replace(f'{{{start:06d}..{end:06d}}}', f'{i:06d}') for i in range(start, end+1)]) |
| 146 | else: |
| 147 | tars.append(shard) |
| 148 | tars = [os.path.join(self.tar_base, t) for t in tars] |
| 149 | # random shuffle the shards |