ADD data loader with multiple tasks in a single mini-batch
@@ -28,7 +28,8 @@ class MultiTaskDataLoader:
:param data_dict: a dictionary with task ids as keys and Datasets
as values.
- :param kwargs: data loader arguments. See pytorch :class:`DataLoader`.
+ :param kwargs: data loader arguments used to instantiate the loader for
+ each task separately. See pytorch :class:`DataLoader`.
"""
self.data_dict = data_dict
self.loaders_dict: Dict[int, DataLoader] = {}
@@ -62,3 +63,58 @@ class MultiTaskDataLoader:
def __len__(self):
return self.max_len * len(self.loaders_dict)
+
+
+class MultiTaskMultiBatchDataLoader:
+ def __init__(self, data_dict: Dict, **kwargs):
+ """ Custom data loader for multi-task training.
+ The dictionary `data_dict` maps task ids into their
+ corresponding datasets.
+
+ mini-batches emitted by this dataloader are dictionaries with task
+ labels as keys and mini-batches as values. Therefore, each mini-batch
+ contains separate data for each task (i.e. key 1 batch for task 1).
+ If datasets for different tasks have different lengths, smaller tasks
+ are oversampled to match the largest task.
+
+ It is suggested to use this loader only if tasks have approximately the
+ same length.
+
+ :param data_dict: a dictionary with task ids as keys and Datasets
+ as values.
+ :param kwargs: data loader arguments used to instantiate the loader for
+ each task separately. See pytorch :class:`DataLoader`.
+ """
+ self.data_dict = data_dict
+ self.loaders_dict: Dict[int, DataLoader] = {}
+
+ for task_id, data in self.data_dict.items():
+ self.loaders_dict[task_id] = DataLoader(data, **kwargs)
+ self.max_len = max([len(d) for d in self.loaders_dict.values()])
+
+ def __iter__(self):
+ iter_dataloaders = {}
+ for t in self.loaders_dict.keys():
+ iter_dataloaders[t] = iter(self.loaders_dict[t])
+ max_len = max([len(d) for d in iter_dataloaders.values()])
+
+ try:
+ for it in range(max_len):
+ mb_curr = {}
+ for t in self.data_dict.keys():
+ t_loader = iter_dataloaders[t]
+ try:
+ x, y = next(t_loader)
+ except StopIteration:
+ # StopIteration is thrown if dataset ends.
+ # reinitialize data loader
+ iter_dataloaders[t] = iter(t_loader)
+ self.current_dataloader = iter_dataloaders[t]
+ x, y = next(t_loader)
+ mb_curr[t] = x, y
+ yield mb_curr
+ except StopIteration:
+ return
+
+ def __len__(self):
+ return self.max_len
