class_docstr: Custom data loader for rehearsal/replay strategies.
class_docstr_tok: ["Custom","data","loader","for","rehearsal\/replay","strategies","."]
docstr: Custom data loader for rehearsal strategies.
  
  This dataloader iterates in parallel two datasets, the current `data`
  and the rehearsal `memory`, which are used to create mini-batches by
  concatenating their data together. Mini-batches from both of them are
  balanced using the task label (i.e. each mini-batch contains a balanced
  number of examples from all the tasks in the `data` and `memory`).
  
  The length of the loader is determined only by the current
  task data and is the same than what it would be when creating a
  data loader for this dataset.
  
  If `oversample_small_tasks == True` smaller tasks are oversampled to
  match the largest task.
  
  :param data: AvalancheDataset.
  :param memory: AvalancheDataset.
  :param oversample_small_tasks: whether smaller tasks should be
      oversampled to match the largest one.
  :param batch_size: the size of the data batch. It must be greater
      than or equal to the number of tasks.
  :param batch_size_mem: the size of the memory batch. If
      `task_balanced_dataloader` is set to True, it must be greater than
      or equal to the number of tasks.
  :param task_balanced_dataloader: if true, buffer data loaders will be
      task-balanced, otherwise it creates a single data loader for the
      buffer samples.
  :param distributed_sampling: If True, apply the PyTorch
      :class:`DistributedSampler`. Defaults to True.
      Note: the distributed sampler is not applied if not running
      a distributed training, even when True is passed.
  :param kwargs: data loader arguments used to instantiate the loader for
      each task separately. See pytorch :class:`DataLoader`.
docstr_tok: ["Custom","data","loader","for","rehearsal","strategies",".","This","dataloader","iterates","in","parallel","two","datasets",",","the","current","`","data","`","and","the","rehearsal","`","memory","`",",","which","are","used","to","create","mini-batches","by","concatenating","their","data","together",".","Mini-batches","from","both","of","them","are","balanced","using","the","task","label","(","i.e",".","each","mini-batch","contains","a","balanced","number","of","examples","from","all","the","tasks","in","the","`","data","`","and","`","memory","`",")",".","The","length","of","the","loader","is","determined","only","by","the","current","task","data","and","is","the","same","than","what","it","would","be","when","creating","a","data","loader","for","this","dataset",".","If","`","oversample_small_tasks","==","True","`","smaller","tasks","are","oversampled","to","match","the","largest","task",".",":","param","data",":","AvalancheDataset",".",":","param","memory",":","AvalancheDataset",".",":","param","oversample_small_tasks",":","whether","smaller","tasks","should","be","oversampled","to","match","the","largest","one",".",":","param","batch_size",":","the","size","of","the","data","batch",".","It","must","be","greater","than","or","equal","to","the","number","of","tasks",".",":","param","batch_size_mem",":","the","size","of","the","memory","batch",".","If","`","task_balanced_dataloader","`","is","set","to","True",",","it","must","be","greater","than","or","equal","to","the","number","of","tasks",".",":","param","task_balanced_dataloader",":","if","true",",","buffer","data","loaders","will","be","task-balanced",",","otherwise","it","creates","a","single","data","loader","for","the","buffer","samples",".",":","param","distributed_sampling",":","If","True",",","apply","the","PyTorch",":","class",":","`","DistributedSampler","`",".","Defaults","to","True",".","Note",":","the","distributed","sampler","is","not","applied","if","not","running","a","distributed","training",",","even","when","True","is","passed",".",":","param","kwargs",":","data","loader","arguments","used","to","instantiate","the","loader","for","each","task","separately",".","See","pytorch",":","class",":","`","DataLoader","`","."]
code: def __init__(
      self,
      data: AvalancheDataset,
      memory: Optional[AvalancheDataset] = None,
      oversample_small_tasks: bool = False,
      batch_size: int = 32,
      batch_size_mem: int = 32,
      task_balanced_dataloader: bool = False,
      distributed_sampling: bool = True,
      **kwargs,
  ):
      """Custom data loader for rehearsal strategies.
  
      This dataloader iterates in parallel two datasets, the current `data`
      and the rehearsal `memory`, which are used to create mini-batches by
      concatenating their data together. Mini-batches from both of them are
      balanced using the task label (i.e. each mini-batch contains a balanced
      number of examples from all the tasks in the `data` and `memory`).
  
      The length of the loader is determined only by the current
      task data and is the same than what it would be when creating a
      data loader for this dataset.
  
      If `oversample_small_tasks == True` smaller tasks are oversampled to
      match the largest task.
  
      :param data: AvalancheDataset.
      :param memory: AvalancheDataset.
      :param oversample_small_tasks: whether smaller tasks should be
          oversampled to match the largest one.
      :param batch_size: the size of the data batch. It must be greater
          than or equal to the number of tasks.
      :param batch_size_mem: the size of the memory batch. If
          `task_balanced_dataloader` is set to True, it must be greater than
          or equal to the number of tasks.
      :param task_balanced_dataloader: if true, buffer data loaders will be
          task-balanced, otherwise it creates a single data loader for the
          buffer samples.
      :param distributed_sampling: If True, apply the PyTorch
          :class:`DistributedSampler`. Defaults to True.
          Note: the distributed sampler is not applied if not running
          a distributed training, even when True is passed.
      :param kwargs: data loader arguments used to instantiate the loader for
          each task separately. See pytorch :class:`DataLoader`.
      """
  
      if "collate_fn" not in kwargs:
          kwargs["collate_fn"] = data.collate_fn
  
      # Create dataloader for memory items
      if task_balanced_dataloader:
          memory_task_labels = getattr(memory, "targets_task_labels")
          assert isinstance(memory_task_labels, DataAttribute)
          num_keys = len(memory_task_labels.uniques)
  
          # Ensure that the per-task batch size will end up > 0
          assert batch_size_mem >= num_keys, (
              "Batch size must be greator or equal "
              "to the number of tasks in the memory "
              "and current data."
          )
  
          # Make the batch size balanced between tasks
          # The remainder (remaining_example) will be distributed
          # across tasks by "self._get_datasets_and_batch_sizes(...)"
          single_group_batch_size = batch_size_mem // num_keys
          remaining_example = batch_size_mem % num_keys
      else:
          single_group_batch_size = batch_size_mem
          remaining_example = 0
  
      # For current data, use the batch_size from the input "batch_size".
      # batch_size can be an int (do not split by task)
      # or a dictionary task_id -> mb_size
      # In both cases, remaining_examples=0
      (
          data_batch_sizes,
          data_subsets,
      ) = self._get_datasets_and_batch_sizes(data, batch_size, 0, False)
  
      (
          memory_batch_sizes,
          memory_subsets,
      ) = self._get_datasets_and_batch_sizes(
          memory,
          single_group_batch_size,
          remaining_example,
          task_balanced_dataloader,
      )
  
      # Obtain the subset with the highest number of iterations
      # This is the one that defines when an epoch ends
      # Note: this is aligned with the behavior of the legacy
      # multi-loader version of ReplayDataLoader
      loaders_for_len_estimation = []
  
      for data_subset, subset_mb_size in zip(
          data_subsets, data_batch_sizes
      ):
          loaders_for_len_estimation.append(
              _make_data_loader(
                  data_subset,
                  distributed_sampling,
                  kwargs,
                  subset_mb_size,
                  force_no_workers=True,
              )[0]
          )
  
      longest_data_subset_idx = (
          np.array(len(d) for d in loaders_for_len_estimation)
          .argmax()
          .item()
      )
  
      super().__init__(
          data_subsets + memory_subsets,
          data_batch_sizes + memory_batch_sizes,
          termination_dataset=longest_data_subset_idx,
          oversample_small_datasets=oversample_small_tasks,
          distributed_sampling=distributed_sampling,
          **kwargs,
      )
  
imports: from typing import Any, Dict, Iterator, List, Optional, Sequence, Sized, Union
  import numpy
  from torch.utils.data import DistributedSampler, Dataset
  from torch.utils.data.dataloader import DataLoader
  from avalanche.benchmarks.utils.data import AvalancheDataset
  from avalanche.benchmarks.utils.data_attribute import DataAttribute
  from avalanche.benchmarks.utils.ffcv_support.ffcv_components import HybridFfcvLoader, has_ffcv_support
  from avalanche.distributed.distributed_helper import DistributedHelper
  from torch.utils.data.sampler import Sampler, BatchSampler
  from torch.utils.data import ConcatDataset
label_desc: Use image node_id 1 to create a new ReplayDataLoader object from inherited base classes: MultiDatasetDataLoader with example: obj = ReplayDataLoader(data, memory, oversample_small_tasks, batch_size, batch_size_mem, task_balanced_dataloader, distributed_sampling)