File size: 3,414 Bytes
cb2428f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
from typing import Optional
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader
class BatchSamplerShard:
def __init__(self, total_samples: int, batch_size: int, shuffle: bool, drop_last: bool, data_seed: Optional[int]):
self.total_samples = total_samples // self.world_size
self.batch_size = batch_size
self.shuffle = shuffle
self.drop_last = drop_last
self.base_seed = data_seed or 0
self.curr_seed = self.base_seed
@property
def rank(self):
return dist.get_rank() if dist.is_initialized() else 0
@property
def world_size(self):
return dist.get_world_size() if dist.is_initialized() else 1
def __iter__(self):
start_idx = self.rank * self.total_samples
if self.shuffle:
generator = torch.Generator()
generator.manual_seed(self.curr_seed)
total_idx = torch.randperm(self.total_samples * self.world_size, generator=generator).tolist()
total_idx = total_idx[start_idx:start_idx + self.total_samples]
else:
total_idx = list(range(start_idx, start_idx + self.total_samples))
batch = []
# Last batch if not complete will be dropped.
for idx in total_idx:
batch.append(idx)
if len(batch) == self.batch_size:
yield batch
batch = []
if not self.drop_last and len(batch) > 0:
yield batch
return
def set_epoch(self, epoch: int):
self.curr_seed = self.base_seed + epoch
def __len__(self) -> int:
if self.drop_last:
return self.total_samples // self.batch_size
else:
return (self.total_samples + self.batch_size - 1) // self.batch_size
class DataLoaderShard(DataLoader):
def __init__(self, dataset, batch_sampler: BatchSamplerShard, **dataloader_params):
self.batch_sampler = batch_sampler
super().__init__(dataset, batch_sampler=self.batch_sampler, **dataloader_params)
def set_epoch(self, epoch: int):
self.batch_sampler.set_epoch(epoch)
class DataLoaderDispatcher:
def __init__(self, base_dataloader):
self.base_dataloader = base_dataloader
@property
def rank(self):
return dist.get_rank(self.group) if dist.is_initialized() else 0
@property
def world_size(self):
return dist.get_world_size(self.group) if dist.is_initialized() else 1
@property
def group(self):
return dist.group.WORLD if dist.is_initialized() else 1
def _scatter_object_list(self, inputs):
if not dist.is_initialized():
return inputs[0]
outputs = [None]
global_src_rank = dist.get_global_rank(self.group, 0)
dist.scatter_object_list(outputs, inputs, global_src_rank, group=self.group)
return outputs[0]
def __iter__(self):
base_iter = iter(self.base_dataloader)
while True:
if self.rank == 0:
try:
data = [next(base_iter) for _ in range(self.world_size)]
except StopIteration:
data = [None] * self.world_size
data = self._scatter_object_list(data)
else:
data = self._scatter_object_list(None)
if data is None:
break
yield data
|