| 193 | return rank, world_size |
| 194 | |
| 195 | class DistributedGivenIterationSampler(Sampler): |
| 196 | def __init__(self, dataset, total_iter, batch_size, world_size=None, rank=None, last_iter=-1, |
| 197 | shuffle_strategy=0, random_seed=0, imageNumPerClass=4, ret_save_path=None): |
| 198 | if world_size is None: |
| 199 | world_size = get_world_size() |
| 200 | if rank is None: |
| 201 | rank = get_rank() |
| 202 | assert rank < world_size |
| 203 | sync_print('sampler: rank={}, world_size={}, random_seed={}'.format(rank, world_size, random_seed)) |
| 204 | self.dataset = dataset |
| 205 | self.total_iter = total_iter |
| 206 | self.batch_size = batch_size |
| 207 | self.world_size = world_size |
| 208 | self.rank = rank |
| 209 | self.last_iter = last_iter |
| 210 | self.shuffle_strategy = shuffle_strategy |
| 211 | self.random_seed = random_seed |
| 212 | self.imageNumPerClass = imageNumPerClass |
| 213 | self.ret_save_path = ret_save_path |
| 214 | self.task_name = self.dataset.task_name |
| 215 | |
| 216 | self.total_size = self.total_iter*self.batch_size |
| 217 | |
| 218 | self.call = 0 |
| 219 | |
| 220 | # generate indices |
| 221 | if self.ret_save_path is not None: |
| 222 | self.this_ret_path = os.path.join(self.ret_save_path, '_'.join([self.task_name, str(self.world_size), str(self.rank)]) + ".pth.tar") |
| 223 | if os.path.exists(self.this_ret_path): |
| 224 | ret_file = torch.load(self.this_ret_path) |
| 225 | # ensure this task and task size is unchanged |
| 226 | if ret_file['task_name'] == self.task_name and ret_file['task_size'] == self.world_size and ret_file['task_rank'] == self.rank: |
| 227 | printlog(" load task sampler from ------> {}".format(self.this_ret_path)) |
| 228 | self.indices = ret_file['ret_file'] |
| 229 | self.dataset.received_indices = True |
| 230 | return |
| 231 | else: |
| 232 | printlog("sampler file ({}) is not existed, and will be generated now--->".format(self.this_ret_path)) |
| 233 | |
| 234 | if self.shuffle_strategy in [0,1,3,4,6]: |
| 235 | self.indices = self.gen_new_list() |
| 236 | self.dataset.indices = self.indices |
| 237 | self.dataset.received_indices = True |
| 238 | elif self.shuffle_strategy == 2: |
| 239 | self.indices = self.gen_s2() |
| 240 | elif self.shuffle_strategy == 5: |
| 241 | self.indices = self.gen_s5() |
| 242 | else: |
| 243 | raise Error("Invalid shuffle_strategy!") # todo: undefined 'Error'??? |
| 244 | |
| 245 | if self.ret_save_path is not None and not os.path.exists(self.ret_save_path): |
| 246 | self.save() |
| 247 | |
| 248 | def gen_s2(self): |
| 249 | |
| 250 | np.random.seed(self.rank) # set different random seed |
| 251 | |
| 252 | indices = [] |
no outgoing calls
no test coverage detected