def __iter__(self):
"""
To guarantee all ranks have the same same permutation, generating it from rank 0 and sync
to other rank by writing to disk
"""
if get_local_rank() == 0:
np.save(self._SAMPLE_FILE, np.array(super().__iter__()))
torch.distributed.barrier()
sample = np.load(self._SAMPLE_FILE)
return iter(sample)
def _dist_permutation(size):
"""Generate permutation for dataset shuffle
Args:
size (int): Size and high value of permutation
Returns:
permutation (ndarray):
"""
if dist.get_world_size() > 1:
# To guarantee all ranks have the same same permutation, generating it from rank 0 and sync
# to other rank by writing to disk
permutation_file = "/tmp/permutation.npy"
if dist.get_local_rank() == 0:
np.save(permutation_file, np.random.permutation(size))
torch.distributed.barrier()
permutation = np.load(permutation_file)
else:
permutation = np.random.permutation(size)
return permutation