Repository: continualai/avalanche
Issue: LwF plugin fails in combination with IncrementalClassifier

🐛 **Describe the bug**
LwFPlugin does not work when combined with IncrementalClassifier.

The error raised is the following (raised when starting to train on second experience):
in lwf.py - _distillation_loss -> line 45
The size of tensor a (6) must match the size of tensor b (10) at non-singleton dimension 1

I believe this is due to the fact that the size of the previous model output and the current model output are different.

I initially thought that this occurred because the option class_ids_from_zero_from_first_exp was not activated, which makes the IncrementalClassifier head as big as the maximum class value (I think it should be documented that IncrementalClassifier should be used in conjuction with this option). But in the end, even when turning this option on, the bug was still appearing but this time with a logical increment between the two output sizes (prev_out.size(0) = out.size(0) - n_classes_in_exp).

🐜 **To Reproduce**
In the example examples/lwf_mnist.py, change the model classifier to an IncrementalClassifier

```{python}
model = SimpleMLP(hidden_size=args.hs)
model.classifier = IncrementalClassifier(model.classifier.in_features)
```

🐝 **Expected behavior**
I think it should be possible to combine this plugin with this classifier type, either by masking the new logits when computing the loss, or by continuing to adapt the head of the old model. At least an error should be raised if the combination is not deemed possible.



Changed files:
avalanche/models/dynamic_modules.py, avalanche/models/pnn.py, avalanche/training/strategies/strategy_wrappers.py, avalanche/training/plugins/ewc.py, tests/training/test_strategies.py, avalanche/training/plugins/lwf.py

diff --git a/avalanche/models/dynamic_modules.py b/avalanche/models/dynamic_modules.py
index b0ca174d..1486a937 100644
--- a/avalanche/models/dynamic_modules.py
+++ b/avalanche/models/dynamic_modules.py
@@ -77,7 +77,7 @@ class DynamicModule(Module):
         pass
 
 
-class MultiTaskModule(Module):
+class MultiTaskModule(DynamicModule):
     """
         Multi-task modules are `torch.nn.Modules`s for multi-task
         scenarios. The `forward` method accepts task labels, one for
@@ -86,7 +86,23 @@ class MultiTaskModule(Module):
         By default the `forward` method splits the mini-batch by task
         and calls `forward_single_task`. Subclasses must implement
         `forward_single_task` or override `forward.
+
+        if `task_labels == None`, the output is computed in parallel
+        for each task.
     """
+    def __init__(self):
+        super().__init__()
+        self.known_train_tasks_labels = set()
+        """ Set of task labels encountered up to now. """
+
+    def train_adaptation(self, dataset: AvalancheDataset = None):
+        """ Update known task labels. """
+        task_labels = dataset.targets_task_labels
+        if isinstance(task_labels, ConstantSequence):
+            # task label is unique. Don't check duplicates.
+            task_labels = [task_labels[0]]
+        self.known_train_tasks_labels = \\
+            self.known_train_tasks_labels.union(set(task_labels))
 
     def forward(self, x: torch.Tensor, task_labels: torch.Tensor)\\
             -> torch.Tensor:
@@ -132,13 +148,17 @@ class MultiTaskModule(Module):
 
     def forward_all_tasks(self, x: torch.Tensor):
         """ compute the output given the input `x` and task label.
+        By default, it considers only tasks seen at training time.
 
         :param x:
         :return: all the possible outputs are returned as a dictionary
             with task IDs as keys and the output of the corresponding
             task as output.
         """
-        raise NotImplementedError()
+        res = {}
+        for task_id in self.known_train_tasks_labels:
+            res[task_id] = self.forward_single_task(x, task_id)
+        return res
 
 
 class IncrementalClassifier(DynamicModule):
@@ -185,7 +205,7 @@ class IncrementalClassifier(DynamicModule):
         return self.classifier(x)
 
 
-class MultiHeadClassifier(MultiTaskModule, DynamicModule):
+class MultiHeadClassifier(MultiTaskModule):
     def __init__(self, in_features, initial_out_features=2):
         """ Multi-head classifier with separate heads for each task.
 
@@ -229,6 +249,7 @@ class MultiHeadClassifier(MultiTaskModule, DynamicModule):
         :param dataset: data from the current experience.
         :return:
         """
+        super().adaptation(dataset)
         task_labels = dataset.targets_task_labels
         if isinstance(task_labels, ConstantSequence):
             # task label is unique. Don't check duplicates.
@@ -252,12 +273,6 @@ class MultiHeadClassifier(MultiTaskModule, DynamicModule):
         """
         return self.classifiers[str(task_label)](x)
 
-    def forward_all_tasks(self, x: torch.Tensor):
-        res = {}
-        for task_id in self.classifiers:
-            res[task_id] = self.classifiers[task_id](x)
-        return res
-
 
 class TrainEvalModel(DynamicModule):
     """
diff --git a/avalanche/models/pnn.py b/avalanche/models/pnn.py
index 09a8bf7a..5315e1ac 100644
--- a/avalanche/models/pnn.py
+++ b/avalanche/models/pnn.py
@@ -107,7 +107,7 @@ class PNNColumn(nn.Module):
         return hs
 
 
-class PNNLayer(MultiTaskModule, DynamicModule):
+class PNNLayer(MultiTaskModule):
     def __init__(self, in_features, out_features_per_column, adapter='mlp'):
         """ Progressive Neural Network layer.
 
@@ -143,6 +143,7 @@ class PNNLayer(MultiTaskModule, DynamicModule):
         :param dataset:
         :return:
         """
+        super().train_adaptation(dataset)
         task_labels = dataset.targets_task_labels
         if isinstance(task_labels, ConstantSequence):
             # task label is unique. Don't check duplicates.
diff --git a/avalanche/training/plugins/ewc.py b/avalanche/training/plugins/ewc.py
index 20b7eb76..12b0b24f 100644
--- a/avalanche/training/plugins/ewc.py
+++ b/avalanche/training/plugins/ewc.py
@@ -40,7 +40,10 @@ class EWCPlugin(StrategyPlugin):
         """
 
         super().__init__()
-
+        assert (decay_factor is None) or (mode == 'online'), \\
+            "You need to set `online` mode to use `decay_factor`."
+        assert (decay_factor is not None) or (mode != 'online'), \\
+            "You need to set `decay_factor` to use the `online` mode."
         assert mode == 'separate' or mode == 'online', \\
             'Mode must be separate or online.'
 
diff --git a/avalanche/training/plugins/lwf.py b/avalanche/training/plugins/lwf.py
index 150f69fd..ffb60bc0 100644
--- a/avalanche/training/plugins/lwf.py
+++ b/avalanche/training/plugins/lwf.py
@@ -39,12 +39,11 @@ class LwFPlugin(StrategyPlugin):
         Compute distillation loss between output of the current model and
         and output of the previous (saved) model.
         """
-
-        log_p = torch.log_softmax(out / self.temperature, dim=1)
-        q = torch.softmax(prev_out / self.temperature, dim=1)
-        res = torch.nn.functional.kl_div(log_p, q, reduction='none')
-        res = res[:, list(active_units)]  # mask unused units
-        res = res.sum() / out.shape[0]  # batch-mean
+        # we compute the loss only on the previously active units.
+        au = list(active_units)
+        log_p = torch.log_softmax(out / self.temperature, dim=1)[:, au]
+        q = torch.softmax(prev_out / self.temperature, dim=1)[:, au]
+        res = torch.nn.functional.kl_div(log_p, q, reduction='batchmean')
         return res
 
     def penalty(self, out, x, alpha, curr_model):
diff --git a/avalanche/training/strategies/strategy_wrappers.py b/avalanche/training/strategies/strategy_wrappers.py
index f5688c11..3af0838b 100644
--- a/avalanche/training/strategies/strategy_wrappers.py
+++ b/avalanche/training/strategies/strategy_wrappers.py
@@ -493,7 +493,6 @@ class EWC(BaseStrategy):
                 if >0: calls `eval` every `eval_every` epochs and at the end
                     of all the epochs for a single experience.
         """
-
         ewc = EWCPlugin(ewc_lambda, mode, decay_factor, keep_importance_data)
         if plugins is None:
             plugins = [ewc]
diff --git a/tests/training/test_strategies.py b/tests/training/test_strategies.py
index cfb7596d..b658141c 100644
--- a/tests/training/test_strategies.py
+++ b/tests/training/test_strategies.py
@@ -18,7 +18,7 @@ from torch.optim import SGD
 from torch.nn import CrossEntropyLoss, Linear
 
 from avalanche.logging import TextLogger
-from avalanche.models import SimpleMLP
+from avalanche.models import SimpleMLP, IncrementalClassifier
 from avalanche.training.plugins import EvaluationPlugin, StrategyPlugin, \\
     LwFPlugin, ReplayPlugin
 from avalanche.training.strategies import Naive, Replay, CWRStar, \\
@@ -37,6 +37,7 @@ from tests.unit_tests_utils import get_fast_benchmark, get_device
 class BaseStrategyTest(unittest.TestCase):
     def test_periodic_eval(self):
         model = SimpleMLP(input_size=6, hidden_size=10)
+        model.classifier = IncrementalClassifier(model.classifier.in_features)
         benchmark = get_fast_benchmark()
         optimizer = SGD(model.parameters(), lr=1e-3)
         criterion = CrossEntropyLoss()
@@ -458,9 +459,15 @@ class StrategyTest(unittest.TestCase):
 
     def get_model(self, fast_test=False):
         if fast_test:
-            return SimpleMLP(input_size=6, hidden_size=10)
+            model = SimpleMLP(input_size=6, hidden_size=10)
+            # model.classifier = IncrementalClassifier(
+            #     model.classifier.in_features)
+            return model
         else:
-            return SimpleMLP()
+            model = SimpleMLP()
+            # model.classifier = IncrementalClassifier(
+            #     model.classifier.in_features)
+            return model
 
     def run_strategy(self, benchmark, cl_strategy):
         print('Starting experiment...')