| 139 | |
| 140 | |
| 141 | def get_train_loader(engine, dataset, config): |
| 142 | data_setting = { |
| 143 | "rgb_root": config.rgb_root_folder, |
| 144 | "rgb_format": config.rgb_format, |
| 145 | "gt_root": config.gt_root_folder, |
| 146 | "gt_format": config.gt_format, |
| 147 | "transform_gt": config.gt_transform, |
| 148 | "x_root": config.x_root_folder, |
| 149 | "x_format": config.x_format, |
| 150 | "x_single_channel": config.x_is_single_channel, |
| 151 | "class_names": config.class_names, |
| 152 | "train_source": config.train_source, |
| 153 | "eval_source": config.eval_source, |
| 154 | "class_names": config.class_names, |
| 155 | "dataset_name": config.dataset_name, |
| 156 | "backbone": config.backbone, |
| 157 | } |
| 158 | train_preprocess = TrainPre(config.norm_mean, config.norm_std, config.x_is_single_channel, config) |
| 159 | |
| 160 | train_dataset = dataset( |
| 161 | data_setting, |
| 162 | "train", |
| 163 | train_preprocess, |
| 164 | config.batch_size * config.niters_per_epoch, |
| 165 | ) |
| 166 | |
| 167 | train_sampler = None |
| 168 | is_shuffle = True |
| 169 | batch_size = config.batch_size |
| 170 | |
| 171 | if engine.distributed: |
| 172 | train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) |
| 173 | batch_size = config.batch_size // engine.world_size |
| 174 | is_shuffle = False |
| 175 | |
| 176 | train_loader = data.DataLoader( |
| 177 | train_dataset, |
| 178 | batch_size=batch_size, |
| 179 | num_workers=config.num_workers, |
| 180 | drop_last=True, |
| 181 | shuffle=is_shuffle, |
| 182 | pin_memory=True, |
| 183 | sampler=train_sampler, |
| 184 | # worker_init_fn=seed_worker, |
| 185 | # generator=g, |
| 186 | ) |
| 187 | |
| 188 | return train_loader, train_sampler |
| 189 | |
| 190 | |
| 191 | def get_val_loader(engine, dataset, config, val_batch_size=1): |