Add possibility to use task-specific transformations.

Author: arthurdouillard

continuum/scenarios/base.py

@@ -16,15 +16,16 @@ class _BaseScenario(abc.ABC):
 
     :param cl_dataset: A Continuum dataset.
     :param nb_tasks: The number of tasks to do.
-    :param transformations: The PyTorch transformations.
-    :param train: Boolean flag whether to use the train or test subset.
+    :param transformations: A list of transformations applied to all tasks. If
+                            it's a list of list, then the transformation will be
+                            different per task.
     """
 
     def __init__(
             self,
             cl_dataset: _ContinuumDataset,
             nb_tasks: int,
-            transformations: List[Callable] = None
+            transformations: Union[List[Callable], List[List[Callable]]] = None
     ) -> None:
 
         self.cl_dataset = cl_dataset
@@ -33,9 +34,15 @@ def __init__(
         if transformations is None:
             transformations = self.cl_dataset.transformations
         if self.cl_dataset.data_type == "segmentation":
-            self.trsf = SegmentationCompose(transformations)
+            composer = SegmentationCompose
         else:
-            self.trsf = transforms.Compose(transformations)
+            composer = transforms.Compose
+        if transformations is not None and isinstance(transformations[0], list):
+            # We have list of list of callable, where each sublist is dedicated to
+            # a task.
+            self.trsf = [composer(trsf) for trsf in transformations]
+        else:
+            self.trsf = composer(transformations)
 
     @abc.abstractmethod
     def _setup(self, nb_tasks: int) -> int:
@@ -87,10 +94,17 @@ def __getitem__(self, task_index: Union[int, slice]):
                            even slices.
         :return: A train PyTorch's Datasets.
         """
+        if isinstance(task_index, slice) and isinstance(self.trsf, list):
+            raise ValueError(
+                f"You cannot select multiple task ({task_index}) when you have a "
+                "different set of transformations per task"
+            )
+
         x, y, t, _ = self._select_data_by_task(task_index)
+
         return TaskSet(
             x, y, t,
-            trsf=self.trsf,
+            trsf=self.trsf[task_index] if isinstance(self.trsf, list) else self.trsf,
             data_type=self.cl_dataset.data_type,
             bounding_boxes=self.cl_dataset.bounding_boxes
         )

continuum/scenarios/class_incremental.py

@@ -20,7 +20,9 @@ class ClassIncremental(_BaseScenario):
                      (e.g. increment=[5,1,1,1,1]).
     :param initial_increment: A different task size applied only for the first task.
                               Desactivated if `increment` is a list.
-    :param transformations: A list of transformations applied to all tasks.
+    :param transformations: A list of transformations applied to all tasks. If
+                            it's a list of list, then the transformation will be
+                            different per task.
     :param class_order: An optional custom class order, used for NC.
                         e.g. [0,1,2,3,4,5,6,7,8,9] or [5,2,4,1,8,6,7,9,0,3]
     """
@@ -31,7 +33,7 @@ def __init__(
         nb_tasks: int = 0,
         increment: Union[List[int], int] = 0,
         initial_increment: int = 0,
-        transformations: List[Callable] = None,
+        transformations: Union[List[Callable], List[List[Callable]]] = None,
         class_order: Union[List[int], None]=None
     ) -> None:
 

continuum/scenarios/continual_scenario.py

@@ -14,16 +14,18 @@ class ContinualScenario(_BaseScenario):
     Scenario: the scenario is entirely defined by the task label vector in the cl_dataset
 
     :param cl_dataset: A continual dataset.
-    :param transformations: A list of transformations applied to all tasks.
+    :param transformations: A list of transformations applied to all tasks. If
+                            it's a list of list, then the transformation will be
+                            different per task.
     """
 
     def __init__(
             self,
             cl_dataset: _ContinuumDataset,
-            transformations: List[Callable] = None,
+            transformations: Union[List[Callable], List[List[Callable]]] = None,
     ) -> None:
         self.check_data(cl_dataset)
-        super().__init__(cl_dataset=cl_dataset, nb_tasks=self.nb_tasks,transformations=transformations)
+        super().__init__(cl_dataset=cl_dataset, nb_tasks=self.nb_tasks, transformations=transformations)
 
     def check_data(self, cl_dataset: _ContinuumDataset):
         x, y, t = cl_dataset.get_data()
@@ -47,4 +49,4 @@ def check_data(self, cl_dataset: _ContinuumDataset):
 
     #nothing to do in the setup function
     def _setup(self, nb_tasks: int) -> int:
-        return nb_tasks
+        return nb_tasks

continuum/scenarios/instance_incremental.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Union
 
 import numpy as np
 
@@ -14,15 +14,17 @@ class InstanceIncremental(_BaseScenario):
 
     :param cl_dataset: A continual dataset.
     :param nb_tasks: The scenario number of tasks.
-    :param transformations: List of transformations to apply to all tasks.
+    :param transformations: A list of transformations applied to all tasks. If
+                            it's a list of list, then the transformation will be
+                            different per task.
     :param random_seed: A random seed to init random processes.
     """
 
     def __init__(
         self,
         cl_dataset: _ContinuumDataset,
         nb_tasks: Optional[int] = None,
-        transformations: List[Callable] = None,
+        transformations: Union[List[Callable], List[List[Callable]]] = None,
         random_seed: int = 1
     ):
         super().__init__(cl_dataset=cl_dataset, nb_tasks=nb_tasks, transformations=transformations)