Used with BatchDataset; Suppose len_buffer == 5, num_buffers == 6, #GPUs. == 3, then. | buffer {i} | buffer {i+1} ------ | ------------------- | ------------------- rank 0 | 0, 1, 2, 3, 4, | 5, 6, 7, 8, 9 rank 1 | 10, 11, 12, 13, 14, | 15, 16, 17,
| 319 | |
| 320 | |
| 321 | class BatchDistributedSampler(DistributedSampler): |
| 322 | """Used with BatchDataset; Suppose len_buffer == 5, num_buffers == 6, |
| 323 | |
| 324 | #GPUs. |
| 325 | |
| 326 | == 3, then. |
| 327 | |
| 328 | | buffer {i} | buffer {i+1} |
| 329 | ------ | ------------------- | ------------------- |
| 330 | rank 0 | 0, 1, 2, 3, 4, | 5, 6, 7, 8, 9 |
| 331 | rank 1 | 10, 11, 12, 13, 14, | 15, 16, 17, 18, 19 |
| 332 | rank 2 | 20, 21, 22, 23, 24, | 25, 26, 27, 28, 29 |
| 333 | """ |
| 334 | |
| 335 | def __init__(self, dataset: Dataset, **kwargs): |
| 336 | super().__init__(dataset, **kwargs) |
| 337 | self.start_index = 0 |
| 338 | |
| 339 | def __iter__(self): |
| 340 | num_buffers = self.dataset.num_buffers |
| 341 | len_buffer = self.dataset.len_buffer |
| 342 | num_buffers_i = num_buffers // self.num_replicas |
| 343 | num_samples_i = len_buffer * num_buffers_i |
| 344 | |
| 345 | indices_i = ( |
| 346 | np.arange(self.start_index, num_samples_i) + |
| 347 | self.rank * num_samples_i) |
| 348 | indices_i = indices_i.tolist() |
| 349 | |
| 350 | return iter(indices_i) |
| 351 | |
| 352 | def reset(self): |
| 353 | self.start_index = 0 |
| 354 | |
| 355 | def state_dict(self, step) -> dict: |
| 356 | return {'start_index': step} |
| 357 | |
| 358 | def load_state_dict(self, state_dict: dict): |
| 359 | self.start_index = state_dict['start_index'] + 1 |
nothing calls this directly
no outgoing calls
no test coverage detected