Add oversampler.

continuum/tasks/utils.py

@@ -1,13 +1,37 @@
 from typing import Tuple, List
 
-
 import torch
 import numpy as np
 
-from continuum.tasks.base import BaseTaskSet
+from continuum.tasks.base import BaseTaskSet, TaskType
 from continuum.tasks.task_set import TaskSet
 
 
+def get_balanced_sampler(taskset, log=False):
+    """Create a sampler that will balance the dataset.
+
+    You should give the returned sampler to the dataloader with the argument `sampler`.
+
+    :param taskset: A pytorch dataset that implement the TaskSet interface.
+    :param log: Use a log weights. If enabled, there will still be imbalance but
+                on the other hand, the oversampling/downsampling won't be as violent.
+    :return: A PyTorch sampler.
+    """
+    if taskset.data_type in (TaskType.SEGMENTATION, TaskType.OBJ_DETECTION, TaskType.TEXT):
+        raise NotImplementedError(
+            "Samplers are not yet available for the "
+            f"{taskset.data_type} type."
+        )
+
+    y = taskset.get_raw_samples()[1]
+    nb_per_class = np.bincount(y)
+    weights_per_class = 1 / nb_per_class
+    if log:
+        weights_per_class = np.log(weights_per_class)
+
+    weights = weights_per_class[y]
+
+    return torch.utils.data.sampler.WeightedRandomSampler(weights, len(taskset))
 
 
 def split_train_val(dataset: BaseTaskSet, val_split: float = 0.1) -> Tuple[BaseTaskSet, BaseTaskSet]: