ReplayDataloader, can fix new data batch size
@@ -218,6 +218,7 @@ class ReplayDataLoader:
oversample_small_tasks: bool = False,
collate_mbatches=_default_collate_mbatches_fn,
batch_size: int = 32,
+ force_data_batch_size: int = None,
**kwargs):
""" Custom data loader for rehearsal strategies.
@@ -239,6 +240,7 @@ class ReplayDataLoader:
combine the mini-batches obtained separately from each task.
:param batch_size: the size of the batch. It must be greater than or
equal to the number of tasks.
+ :param ratio_data_mem: How many of the samples should be from
:param kwargs: data loader arguments used to instantiate the loader for
each task separately. See pytorch :class:`DataLoader`.
"""
@@ -250,6 +252,24 @@ class ReplayDataLoader:
self.oversample_small_tasks = oversample_small_tasks
self.collate_mbatches = collate_mbatches
+ if force_data_batch_size is not None:
+ assert force_data_batch_size <= batch_size, \
+ "Forced batch size of data must be <= entire batch size"
+
+ mem_batch_size = batch_size - force_data_batch_size
+ remaining_example = 0
+ mem_keys = len(self.memory.task_set)
+ assert mem_batch_size >= mem_keys, \
+ "Batch size must be greator or equal " \
+ "to the number of tasks in the memory."
+
+ self.loader_data, _ = self._create_dataloaders(
+ data, force_data_batch_size,
+ remaining_example, **kwargs)
+ self.loader_memory, _ = self._create_dataloaders(
+ memory, mem_batch_size,
+ remaining_example, **kwargs)
+ else:
num_keys = len(self.data.task_set) + len(self.memory.task_set)
assert batch_size >= num_keys, "Batch size must be greator or equal "\
"to the number of tasks in the memory " \
@@ -264,6 +284,7 @@ class ReplayDataLoader:
self.loader_memory, remaining_example = self._create_dataloaders(
memory, single_group_batch_size,
remaining_example, **kwargs)
+
self.max_len = max([len(d) for d in chain(
self.loader_data.values(), self.loader_memory.values())]
)
