Repository: ContinualAI/avalanche

Issue:
Update LwF with new BaseStrategy
update `training.deprecated.lwf`.


Discussion:

Resolving patch:
diff --git a/avalanche/training/plugins.py b/avalanche/training/plugins.py
--- a/avalanche/training/plugins.py
+++ b/avalanche/training/plugins.py
@@ -677,5 +677,67 @@ def expand_task_layer(self, strategy, min_n_output_units: int, task_layer):
         return new_layer
 
 
+class LwFPlugin(StrategyPlugin):
+    """
+    A Learning without Forgetting plugin.
+    LwF uses distillation to regularize the current loss with soft targets
+    taken from a previous version of the model. 
+    """
+
+    def __init__(self, alpha=1, temperature=2):
+        """
+        :param alpha: distillation hyperparameter. It can be either a float
+                number or a list containing alpha for each step.
+        :param temperature: softmax temperature for distillation
+        """
+
+        super().__init__()
+
+        self.alpha = alpha
+        self.temperature = temperature
+        self.prev_model = None
+        self.step_id = 0
+
+    def _distillation_loss(self, out, prev_out):
+        """
+        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='batchmean')
+        return res
+
+    def penalty(self, out, x, alpha):
+        """
+        Compute weighted distillation loss.
+        """
+
+        if self.prev_model is None:
+            return 0
+        else:
+            y_prev = self.prev_model(x).detach()
+            dist_loss = self._distillation_loss(out, y_prev)
+            return alpha * dist_loss
+
+    def before_backward(self, strategy, **kwargs):
+        """
+        Add distillation loss
+        """
+        alpha = self.alpha[self.step_id] \
+            if isinstance(self.alpha, (list, tuple)) else self.alpha
+        penalty = self.penalty(strategy.logits, strategy.mb_x, alpha)
+        strategy.loss += penalty
+
+    def after_training_step(self, strategy, **kwargs):
+        """
+        Save a copy of the model after each step
+        """
+
+        self.prev_model = copy.deepcopy(strategy.model)
+        self.step_id += 1
+
+
 __all__ = ['StrategyPlugin', 'ReplayPlugin', 'GDumbPlugin',
-           'EvaluationPlugin', 'CWRStarPlugin', 'MultiHeadPlugin']
+           'EvaluationPlugin', 'CWRStarPlugin', 'MultiHeadPlugin', 'LwFPlugin']
diff --git a/avalanche/training/strategies/joint_training.py b/avalanche/training/strategies/joint_training.py
--- a/avalanche/training/strategies/joint_training.py
+++ b/avalanche/training/strategies/joint_training.py
@@ -9,18 +9,16 @@
 # Website: clair.continualai.org                                               #
 ################################################################################
 
-from collections import defaultdict
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence
 
 from torch.nn import Module
 from torch.optim import Optimizer
-from torch.utils.data import DataLoader, ConcatDataset
+from torch.utils.data import ConcatDataset
 
 from avalanche.benchmarks.scenarios import IStepInfo
 from avalanche.training.plugins import StrategyPlugin
 from avalanche.evaluation.eval_protocol import EvalProtocol
 from avalanche.training.strategies import BaseStrategy
-from avalanche.benchmarks.utils import concat_datasets_sequentially
 
 
 class JointTraining(BaseStrategy):
@@ -78,7 +76,7 @@ def train(self, step_infos: Sequence[IStepInfo], **kwargs):
             self.current_data = ConcatDataset(datasets)
             self.step_info = step_infos[0]
         else:
-            raise NotImplemented
+            raise NotImplementedError
 
         print("starting training...")
         self.model.train()
diff --git a/avalanche/training/strategies/strategies.py b/avalanche/training/strategies/strategies.py
--- a/avalanche/training/strategies/strategies.py
+++ b/avalanche/training/strategies/strategies.py
@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, List
+from typing import Optional, Sequence, List, Union
 
 from torch.nn import Module
 from torch.optim import Optimizer
@@ -8,7 +8,7 @@
 from avalanche.evaluation import EvalProtocol
 from avalanche.training.strategies.base_strategy import BaseStrategy
 from avalanche.training.plugins import StrategyPlugin, \
-    CWRStarPlugin, ReplayPlugin, GDumbPlugin
+    CWRStarPlugin, ReplayPlugin, GDumbPlugin, LwFPlugin
 
 
 class Naive(BaseStrategy):
@@ -188,4 +188,42 @@ def adapt_train_dataset(self, **kwargs):
             self.current_data = self.dataset
 
 
-__all__ = ['Naive', 'CWRStar', 'Replay', 'GDumb', 'Cumulative']
+class LwF(BaseStrategy):
+
+    def __init__(self, model: Module, optimizer: Optimizer, criterion,
+                 alpha: Union[float, Sequence[float]], temperature: float,
+                 evaluation_protocol: Optional[EvalProtocol] = None,
+                 train_mb_size: int = 1, train_epochs: int = 1,
+                 test_mb_size: int = None, device=None,
+                 plugins: Optional[List[StrategyPlugin]] = None,
+                 ):
+        """ Learning without Forgetting strategy. 
+            See LwF plugin for details.
+
+        :param model: The model.
+        :param optimizer: The optimizer to use.
+        :param criterion: The loss criterion to use.
+        :param alpha: distillation hyperparameter. It can be either a float
+                number or a list containing alpha for each step.
+        :param temperature: softmax temperature for distillation
+        :param evaluation_protocol: The evaluation plugin.
+        :param train_mb_size: The train minibatch size. Defaults to 1.
+        :param train_epochs: The number of training epochs. Defaults to 1.
+        :param test_mb_size: The test minibatch size. Defaults to 1.
+        :param device: The device to use. Defaults to None (cpu).
+        :param plugins: Plugins to be added. Defaults to None.
+        """
+
+        lwf = LwFPlugin(alpha, temperature)
+        if plugins is None:
+            plugins = [lwf]
+        else:
+            plugins.append(lwf)
+
+        super().__init__(
+            model, optimizer, criterion, evaluation_protocol,
+            train_mb_size=train_mb_size, train_epochs=train_epochs,
+            test_mb_size=test_mb_size, device=device, plugins=plugins)
+
+
+__all__ = ['Naive', 'CWRStar', 'Replay', 'GDumb', 'Cumulative', 'LwF']


Test patch:
diff --git a/tests/test_strategies.py b/tests/test_strategies.py
--- a/tests/test_strategies.py
+++ b/tests/test_strategies.py
@@ -23,7 +23,7 @@
 
 from avalanche.extras.models import SimpleMLP
 from avalanche.training.strategies import Naive, Replay, CWRStar, \
-    GDumb, Cumulative
+    GDumb, Cumulative, LwF
 from avalanche.benchmarks import nc_scenario, NCStepInfo
 
 
@@ -116,6 +116,19 @@ def test_cumulative(self):
 
         strategy = Cumulative(model, optimizer, criterion, train_mb_size=64)
         self.run_strategy(my_nc_scenario, strategy)
+    
+    def test_lwf(self):
+        model = SimpleMLP()
+        optimizer = SGD(model.parameters(), lr=1e-3)
+        criterion = CrossEntropyLoss()
+        mnist_train, mnist_test = self.load_dataset()
+        my_nc_scenario = nc_scenario(
+            mnist_train, mnist_test, 5, task_labels=False, seed=1234)
+
+        strategy = LwF(model, optimizer, criterion,
+                       alpha=[0, 1/2, 2*(2/3), 3*(3/4), 4*(4/5)], 
+                       temperature=2, train_mb_size=64)
+        self.run_strategy(my_nc_scenario, strategy)
 
     def load_dataset(self):
         mnist_train = MNIST(