Repository: continualai/avalanche
Issue: Fix type hints to make use of new benchmark structure

I think we should update the typing in function signatures such as https://github.com/ContinualAI/avalanche/blob/master/avalanche/training/templates/base.py#L74 requiring specific experiences to be more general, replacing them with the new [CLExperience](https://github.com/ContinualAI/avalanche/blob/master/avalanche/benchmarks/scenarios/generic_scenario.py#L73), perhaps even swapping the `Sequence` of experiences requirement with a [CLStream](https://github.com/ContinualAI/avalanche/blob/master/avalanche/benchmarks/scenarios/generic_scenario.py#L154).


Changed files:
avalanche/training/templates/base.py, avalanche/training/templates/online_supervised.py, avalanche/training/templates/base_sgd.py

diff --git a/avalanche/training/templates/base.py b/avalanche/training/templates/base.py
index 032f2fc1..7b3a1679 100644
--- a/avalanche/training/templates/base.py
+++ b/avalanche/training/templates/base.py
@@ -1,13 +1,15 @@
 import warnings
-from typing import Sequence, Optional, Union, List
+from typing import Iterable, Sequence, Optional, Union, List
 
 import torch
 from torch.nn import Module
 
-from avalanche.benchmarks import ClassificationExperience
+from avalanche.benchmarks import CLExperience, CLStream
 from avalanche.core import BasePlugin
 from avalanche.training.utils import trigger_plugins
 
+ExpSequence = Iterable[CLExperience]
+
 
 class BaseTemplate:
     """Base class for continual learning skeletons.
@@ -52,15 +54,13 @@ class BaseTemplate:
         ###################################################################
         # State variables. These are updated during the train/eval loops. #
         ###################################################################
-        self.experience: Optional[ClassificationExperience] = None
+        self.experience: Optional[CLExperience] = None
         """ Current experience. """
 
         self.is_training: bool = False
         """ True if the strategy is in training mode. """
 
-        self.current_eval_stream: Optional[
-            Sequence[ClassificationExperience]
-        ] = None
+        self.current_eval_stream: Optional[ExpSequence] = None
         """ Current evaluation stream. """
 
     @property
@@ -70,16 +70,9 @@ class BaseTemplate:
 
     def train(
         self,
-        experiences: Union[
-            ClassificationExperience, Sequence[ClassificationExperience]
-        ],
-        eval_streams: Optional[
-            Sequence[
-                Union[
-                    ClassificationExperience, Sequence[ClassificationExperience]
-                ]
-            ]
-        ] = None,
+        experiences: Union[CLExperience, ExpSequence],
+        eval_streams: Optional[Sequence[Union[CLExperience,
+                                              ExpSequence]]] = None,
         **kwargs,
     ):
         """Training loop.
@@ -90,7 +83,7 @@ class BaseTemplate:
         It returns a dictionary with last recorded value for each metric.
 
         :param experiences: single Experience or sequence.
-        :param eval_streams: list of streams for evaluation.
+        :param eval_streams: sequence of streams for evaluation.
             If None: use training experiences for evaluation.
             Use [] if you do not want to evaluate during training.
         """
@@ -101,7 +94,7 @@ class BaseTemplate:
         self.model.to(self.device)
 
         # Normalize training and eval data.
-        if not isinstance(experiences, Sequence):
+        if not isinstance(experiences, Iterable):
             experiences = [experiences]
         if eval_streams is None:
             eval_streams = [experiences]
@@ -116,16 +109,14 @@ class BaseTemplate:
         self._after_training(**kwargs)
 
     def _train_exp(
-        self, experience: ClassificationExperience, eval_streams, **kwargs
+        self, experience: CLExperience, eval_streams, **kwargs
     ):
         raise NotImplementedError()
 
     @torch.no_grad()
     def eval(
         self,
-        exp_list: Union[
-            ClassificationExperience, Sequence[ClassificationExperience]
-        ],
+        exp_list: Union[CLExperience, CLStream],
         **kwargs,
     ):
         """
@@ -144,7 +135,7 @@ class BaseTemplate:
         self.is_training = False
         self.model.eval()
 
-        if not isinstance(exp_list, Sequence):
+        if not isinstance(exp_list, Iterable):
             exp_list = [exp_list]
         self.current_eval_stream = exp_list
 
diff --git a/avalanche/training/templates/base_sgd.py b/avalanche/training/templates/base_sgd.py
index be0559dd..c361dd82 100644
--- a/avalanche/training/templates/base_sgd.py
+++ b/avalanche/training/templates/base_sgd.py
@@ -1,15 +1,15 @@
-from typing import Sequence, Optional, Union, List
+from typing import Iterable, Sequence, Optional, Union, List
 
 import torch
 from torch.nn import Module
 from torch.optim import Optimizer
 
-from avalanche.benchmarks import ClassificationExperience
+from avalanche.benchmarks import CLExperience, CLStream
 from avalanche.core import BaseSGDPlugin
 from avalanche.training.plugins import SupervisedPlugin, EvaluationPlugin
 from avalanche.training.plugins.clock import Clock
 from avalanche.training.plugins.evaluation import default_evaluator
-from avalanche.training.templates.base import BaseTemplate
+from avalanche.training.templates.base import BaseTemplate, ExpSequence
 
 from typing import TYPE_CHECKING
 
@@ -121,31 +121,17 @@ class BaseSGDTemplate(BaseTemplate):
 
         self._stop_training = False
 
-    def train(
-        self,
-        experiences: Union[
-            ClassificationExperience, Sequence[ClassificationExperience]
-        ],
-        eval_streams: Optional[
-            Sequence[
-                Union[
-                    ClassificationExperience, Sequence[ClassificationExperience]
-                ]
-            ]
-        ] = None,
-        **kwargs,
-    ):
+    def train(self,
+              experiences: Union[CLExperience,
+                                 ExpSequence],
+              eval_streams: Optional[Sequence[Union[CLExperience,
+                                                    ExpSequence]]] = None,
+              **kwargs):
         super().train(experiences, eval_streams, **kwargs)
         return self.evaluator.get_last_metrics()
 
     @torch.no_grad()
-    def eval(
-        self,
-        exp_list: Union[
-            ClassificationExperience, Sequence[ClassificationExperience]
-        ],
-        **kwargs,
-    ):
+    def eval(self, exp_list: Union[CLExperience, CLStream], **kwargs):
         """
         Evaluate the current model on a series of experiences and
         returns the last recorded value for each metric.
@@ -167,7 +153,7 @@ class BaseSGDTemplate(BaseTemplate):
         super()._before_training_exp(**kwargs)
 
     def _train_exp(
-        self, experience: ClassificationExperience, eval_streams=None, **kwargs
+        self, experience: CLExperience, eval_streams=None, **kwargs
     ):
         """Training loop over a single Experience object.
 
@@ -180,7 +166,7 @@ class BaseSGDTemplate(BaseTemplate):
         if eval_streams is None:
             eval_streams = [experience]
         for i, exp in enumerate(eval_streams):
-            if not isinstance(exp, Sequence):
+            if not isinstance(exp, Iterable):
                 eval_streams[i] = [exp]
         for _ in range(self.train_epochs):
             self._before_training_epoch(**kwargs)
diff --git a/avalanche/training/templates/online_supervised.py b/avalanche/training/templates/online_supervised.py
index 609e9247..11eec064 100644
--- a/avalanche/training/templates/online_supervised.py
+++ b/avalanche/training/templates/online_supervised.py
@@ -1,17 +1,18 @@
 import copy
 import warnings
-from typing import Optional, List, Union, Sequence
+from typing import Optional, List, Union, Sequence, Iterable
 
 import torch
 from torch.nn import Module, CrossEntropyLoss
 from torch.optim import Optimizer
 
-from avalanche.benchmarks import ClassificationExperience
+from avalanche.benchmarks import CLExperience
 from avalanche.benchmarks.utils import AvalancheSubset
 from avalanche.models import DynamicModule
 from avalanche.training.plugins import SupervisedPlugin, EvaluationPlugin
 from avalanche.training.plugins.evaluation import default_evaluator
 from avalanche.training.templates.supervised import SupervisedTemplate
+from avalanche.training.templates.base import ExpSequence
 
 
 class SupervisedOnlineTemplate(SupervisedTemplate):
@@ -48,7 +49,8 @@ class SupervisedOnlineTemplate(SupervisedTemplate):
             "Some plugins may not work properly."
         )
 
-    def create_sub_experience_list(self, experience):
+    def create_sub_experience_list(
+            self, experience: CLExperience) -> List[CLExperience]:
         """Creates a list of sub-experiences from an experience.
         It returns a list of experiences, where each experience is
         a subset of the original experience.
@@ -78,20 +80,12 @@ class SupervisedOnlineTemplate(SupervisedTemplate):
 
         return sub_experience_list
 
-    def train(
-        self,
-        experiences: Union[
-            ClassificationExperience, Sequence[ClassificationExperience]
-        ],
-        eval_streams: Optional[
-            Sequence[
-                Union[
-                    ClassificationExperience, Sequence[ClassificationExperience]
-                ]
-            ]
-        ] = None,
-        **kwargs
-    ):
+    def train(self,
+              experiences: Union[CLExperience,
+                                 ExpSequence],
+              eval_streams: Optional[Sequence[Union[CLExperience,
+                                                    ExpSequence]]] = None,
+              **kwargs):
         """Training loop. if experiences is a single element trains on it.
         If it is a sequence, trains the model on each experience in order.
         This is different from joint training on the entire stream.
@@ -112,7 +106,7 @@ class SupervisedOnlineTemplate(SupervisedTemplate):
         self.model.to(self.device)
 
         # Normalize training and eval data.
-        if not isinstance(experiences, Sequence):
+        if not isinstance(experiences, Iterable):
             experiences = [experiences]
         if eval_streams is None:
             eval_streams = [experiences]
@@ -148,7 +142,7 @@ class SupervisedOnlineTemplate(SupervisedTemplate):
 
     def _train_exp(
         self,
-        experience: ClassificationExperience,
+        experience: CLExperience,
         eval_streams=None,
         is_first_sub_exp=False,
         is_last_sub_exp=False,
@@ -172,7 +166,7 @@ class SupervisedOnlineTemplate(SupervisedTemplate):
         if eval_streams is None:
             eval_streams = [experience]
         for i, exp in enumerate(eval_streams):
-            if not isinstance(exp, Sequence):
+            if not isinstance(exp, Iterable):
                 eval_streams[i] = [exp]
 
         # Data Adaptation (e.g. add new samples/data augmentation)