|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from torch.utils.data import DataLoader, RandomSampler |
|
|
from torch.utils.data.distributed import DistributedSampler |
|
|
from deepspeed.accelerator import get_accelerator |
|
|
|
|
|
from deepspeed.runtime.data_pipeline.data_sampling.data_sampler import DeepSpeedDataSampler |
|
|
from deepspeed.runtime.data_pipeline.constants import CURRICULUM_LEARNING, \ |
|
|
DATA_EFFICIENCY, DATA_SAMPLING_NUM_WORKERS |
|
|
from deepspeed.runtime.constants import GRADIENT_ACCUMULATION_STEPS, \ |
|
|
DATA_PARALLEL_GROUP, GLOBAL_RANK |
|
|
|
|
|
|
|
|
class RepeatingLoader: |
|
|
|
|
|
def __init__(self, loader): |
|
|
"""Wraps an iterator to allow for infinite iteration. This is especially useful |
|
|
for DataLoader types that we wish to automatically restart upon completion. |
|
|
|
|
|
Args: |
|
|
loader (iterator): The data loader to repeat. |
|
|
""" |
|
|
self.loader = loader |
|
|
self.data_iter = iter(self.loader) |
|
|
|
|
|
def __iter__(self): |
|
|
return self |
|
|
|
|
|
def __next__(self): |
|
|
try: |
|
|
batch = next(self.data_iter) |
|
|
except StopIteration: |
|
|
self.data_iter = iter(self.loader) |
|
|
batch = next(self.data_iter) |
|
|
return batch |
|
|
|
|
|
|
|
|
class DeepSpeedDataLoader(object): |
|
|
|
|
|
def __init__(self, |
|
|
dataset, |
|
|
batch_size, |
|
|
pin_memory, |
|
|
local_rank, |
|
|
tput_timer, |
|
|
collate_fn=None, |
|
|
num_local_io_workers=None, |
|
|
data_sampler=None, |
|
|
data_parallel_world_size=None, |
|
|
data_parallel_rank=None, |
|
|
dataloader_drop_last=False, |
|
|
deepspeed_dataloader_config={}): |
|
|
self.deepspeed_dataloader_config = deepspeed_dataloader_config |
|
|
self.tput_timer = tput_timer |
|
|
self.batch_size = batch_size |
|
|
self.curriculum_learning_enabled = False |
|
|
if CURRICULUM_LEARNING in deepspeed_dataloader_config: |
|
|
self.curriculum_learning_enabled = deepspeed_dataloader_config[CURRICULUM_LEARNING] |
|
|
|
|
|
if self.curriculum_learning_enabled: |
|
|
data_sampler = DeepSpeedDataSampler(self.deepspeed_dataloader_config[DATA_EFFICIENCY], |
|
|
len(dataset), |
|
|
self.batch_size, |
|
|
data_parallel_rank, |
|
|
data_parallel_world_size, |
|
|
self.deepspeed_dataloader_config[DATA_PARALLEL_GROUP], |
|
|
self.deepspeed_dataloader_config[GRADIENT_ACCUMULATION_STEPS], |
|
|
self.deepspeed_dataloader_config[GLOBAL_RANK], |
|
|
drop_last=dataloader_drop_last) |
|
|
device_count = get_accelerator().device_count() |
|
|
num_local_io_workers = self.deepspeed_dataloader_config[DATA_SAMPLING_NUM_WORKERS] |
|
|
else: |
|
|
if local_rank >= 0: |
|
|
if data_sampler is None: |
|
|
data_sampler = DistributedSampler(dataset=dataset, |
|
|
num_replicas=data_parallel_world_size, |
|
|
rank=data_parallel_rank) |
|
|
device_count = 1 |
|
|
else: |
|
|
if data_sampler is None: |
|
|
data_sampler = RandomSampler(dataset) |
|
|
device_count = get_accelerator().device_count() |
|
|
batch_size *= device_count |
|
|
|
|
|
if num_local_io_workers is None: |
|
|
num_local_io_workers = 2 * device_count |
|
|
|
|
|
self.num_local_io_workers = num_local_io_workers |
|
|
self.data_sampler = data_sampler |
|
|
self.dataset = dataset |
|
|
self.collate_fn = collate_fn |
|
|
self.device_count = device_count |
|
|
self.batch_size = batch_size |
|
|
self.pin_memory = pin_memory |
|
|
self.data = None |
|
|
self.dataloader_drop_last = dataloader_drop_last |
|
|
self.post_process_func = None |
|
|
|
|
|
if self.dataloader_drop_last: |
|
|
self.len = len(self.data_sampler) // self.batch_size |
|
|
else: |
|
|
from math import ceil |
|
|
self.len = ceil(len(self.data_sampler) / self.batch_size) |
|
|
|
|
|
def __iter__(self): |
|
|
self._create_dataloader() |
|
|
return self |
|
|
|
|
|
def __len__(self): |
|
|
return self.len |
|
|
|
|
|
def __next__(self): |
|
|
if self.tput_timer: |
|
|
self.tput_timer.start() |
|
|
if self.curriculum_learning_enabled: |
|
|
data = next(self.data_iterator) |
|
|
if self.post_process_func is not None: |
|
|
data = self.post_process_func(data, self.data_sampler.state_dict()) |
|
|
return data |
|
|
else: |
|
|
return next(self.data) |
|
|
|
|
|
def _create_dataloader(self): |
|
|
if self.curriculum_learning_enabled: |
|
|
if self.collate_fn is None: |
|
|
self.dataloader = DataLoader(self.dataset, |
|
|
pin_memory=self.pin_memory, |
|
|
batch_sampler=self.data_sampler, |
|
|
num_workers=self.num_local_io_workers) |
|
|
else: |
|
|
self.dataloader = DataLoader(self.dataset, |
|
|
pin_memory=self.pin_memory, |
|
|
batch_sampler=self.data_sampler, |
|
|
collate_fn=self.collate_fn, |
|
|
num_workers=self.num_local_io_workers) |
|
|
self.data_iterator = iter(self.dataloader) |
|
|
return self.dataloader |
|
|
else: |
|
|
if self.collate_fn is None: |
|
|
self.dataloader = DataLoader(self.dataset, |
|
|
batch_size=self.batch_size, |
|
|
pin_memory=self.pin_memory, |
|
|
sampler=self.data_sampler, |
|
|
num_workers=self.num_local_io_workers, |
|
|
drop_last=self.dataloader_drop_last) |
|
|
else: |
|
|
self.dataloader = DataLoader(self.dataset, |
|
|
batch_size=self.batch_size, |
|
|
pin_memory=self.pin_memory, |
|
|
sampler=self.data_sampler, |
|
|
collate_fn=self.collate_fn, |
|
|
num_workers=self.num_local_io_workers, |
|
|
drop_last=self.dataloader_drop_last) |
|
|
self.data = (x for x in self.dataloader) |
|
|
|
|
|
return self.dataloader |
|
|
|
|
|
|
|
|
|
|
|
|
