Repository: ContinualAI/avalanche

Issue:
Scenario not Necessary in MultiHead Plugin
Apparently, as discusses in #174, the use of the scenario is not necessary in the MultiHead plugin implementation.
We should get rid of it to be more coherent with the API and avoid the hack of accessing to it through the step_info object.


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
@@ -17,6 +17,7 @@ class StrategyPlugin:
     by the BaseStrategy with an empty function. Subclasses must override
     the callbacks.
     """
+
     def __init__(self):
         pass
 
@@ -99,6 +100,7 @@ class ReplayPlugin(StrategyPlugin):
     the external memory. We assume the training set contains at least 
     :mem_size: data points.
     """
+
     def __init__(self, mem_size=200):
         super().__init__()
 
@@ -143,7 +145,7 @@ def after_training(self, strategy, **kwargs):
             ext_mem = copy.deepcopy(self.rm_add)
         else:
             _, saved_part = random_split(
-                ext_mem, [len(self.rm_add), len(ext_mem)-len(self.rm_add)]
+                ext_mem, [len(self.rm_add), len(ext_mem) - len(self.rm_add)]
             )
             ext_mem = ConcatDataset([saved_part, self.rm_add])
         self.ext_mem = ext_mem
@@ -208,7 +210,7 @@ def adapt_train_dataset(self, strategy, **kwargs):
                             self.ext_mem.tensors[1][j] = target
                             break
                     self.counter[to_remove] -= 1
-                else:   
+                else:
                     # memory not full: add new pattern
                     if self.ext_mem is None:
                         self.ext_mem = TensorDataset(
@@ -237,20 +239,26 @@ class EvaluationPlugin(StrategyPlugin):
     required metrics. The "evaluation_protocol" is usually passed as argument
     from the strategy.
     """
+
     def __init__(self, evaluation_protocol):
         super().__init__()
         self.evaluation_protocol = evaluation_protocol
 
+        # Private state variables
+        self._dataset_size = None
+        self._seen_samples = 0
+        self._total_loss = 0
+        self._average_loss = 0
+
         # Training
-        self._training_dataset_size = None
         self._training_accuracy = None
         self._training_correct_count = 0
         self._training_average_loss = 0
         self._training_total_iterations = 0
         self._train_current_task_id = None
+        self._train_seen_samples = 0
 
         # Test
-        self._test_dataset_size = None
         self._test_average_loss = 0
         self._test_current_task_id = None
         self._test_true_y = None
@@ -265,60 +273,60 @@ def get_test_result(self):
 
     def before_training(self, strategy, **kwargs):
         task_id = strategy.step_info.task_label
-        self._training_dataset_size = None
         self._train_current_task_id = task_id
         self._training_accuracy = None
         self._training_correct_count = 0
-        self._training_average_loss = 0
-        self._training_dataset_size = len(strategy.current_data)
+        self._dataset_size = len(strategy.current_data)
+        self._seen_samples = 0
+        self._total_loss = 0
+        self._average_loss = 0
+        print("[Training on Task {}, Step {}]"
+              .format(self._test_current_task_id,
+                      strategy.step_info.current_step))
 
     def after_training_iteration(self, strategy, **kwargs):
         self._training_total_iterations += 1
-        epoch = strategy.epoch
         iteration = strategy.mb_it
         train_mb_y = strategy.mb_y
         logits = strategy.logits
         loss = strategy.loss
-
-        # TODO: is this a bug?
-        den = ((iteration + 1) * train_mb_y.shape[0] +
-               epoch * self._training_dataset_size)
+        self._seen_samples += train_mb_y.shape[0]
 
         # Accuracy
         _, predicted_labels = torch.max(logits, 1)
-        correct_predictions = torch.eq(predicted_labels, train_mb_y)\
-                                   .sum().item()
+        correct_predictions = torch.eq(predicted_labels, train_mb_y) \
+            .sum().item()
         self._training_correct_count += correct_predictions
-        self._training_accuracy = self._training_correct_count / den
+        self._training_accuracy = self._training_correct_count / \
+            self._seen_samples
 
         # Loss
-        self._training_average_loss += loss.item()
-        self._training_average_loss /= den
+        self._total_loss += loss.item()
+        self._average_loss = self._total_loss / self._seen_samples
 
         # Logging
         if iteration % 100 == 0:
             print(
                 '[Training] ==>>> it: {}, avg. loss: {:.6f}, '
                 'running train acc: {:.3f}'.format(
-                    iteration, self._training_average_loss,
+                    iteration, self._average_loss,
                     self._training_accuracy))
 
             self.evaluation_protocol.update_tb_train(
-                self._training_average_loss, self._training_accuracy,
+                self._average_loss, self._training_accuracy,
                 self._training_total_iterations, torch.unique(train_mb_y),
                 self._train_current_task_id)
 
     def before_test_step(self, strategy, **kwargs):
         step_info = strategy.step_info
-        step_id = strategy.step_id
 
-        task_id = step_info.scenario.test_stream[step_id].task_label
         self._test_protocol_results = dict()
-        self._test_dataset_size = len(strategy.current_data)
-        self._test_current_task_id = task_id
+        self._test_current_task_id = step_info.task_label
         self._test_average_loss = 0
         self._test_true_y = []
         self._test_predicted_y = []
+        self._dataset_size = len(strategy.current_data)
+        self._seen_samples = 0
 
     def after_test_iteration(self, strategy, **kwargs):
         _, predicted_labels = torch.max(strategy.logits, 1)
@@ -327,15 +335,17 @@ def after_test_iteration(self, strategy, **kwargs):
         self._test_average_loss += strategy.loss.item()
 
     def after_test_step(self, strategy, **kwargs):
-        self._test_average_loss /= self._test_dataset_size
+        self._test_average_loss /= self._dataset_size
 
         results = self.evaluation_protocol.get_results(
             self._test_true_y, self._test_predicted_y,
             self._train_current_task_id, self._test_current_task_id)
         acc, accs = results[ACC]
 
-        print("[Evaluation] Task {0}: Avg Loss {1}; Avg Acc {2}"
-              .format(self._test_current_task_id, self._test_average_loss, acc))
+        print("[Evaluation] Task {}, Step {}: Avg Loss {:.6f}; Avg Acc {:.3f}"
+              .format(self._test_current_task_id,
+                      strategy.step_info.current_step,
+                      self._test_average_loss, acc))
 
         self._test_protocol_results[self._test_current_task_id] = \
             (self._test_average_loss, acc, accs, results)
@@ -374,7 +384,6 @@ def __init__(self, model, second_last_layer_name, num_classes=50):
     def after_training(self, strategy, **kwargs):
         CWRStarPlugin.consolidate_weights(self.model, self.cur_class)
         self.batch_processed += 1
-        strategy.optimizer
 
     def before_training(self, strategy, **kwargs):
         if self.batch_processed == 1:
@@ -411,7 +420,7 @@ def consolidate_weights(model, cur_clas):
                     if c in model.saved_weights.keys():
                         wpast_j = np.sqrt(model.past_j[c] / model.cur_j[c])
                         # wpast_j = model.past_j[c] / model.cur_j[c]
-                        model.saved_weights[c] = (model.saved_weights[c] * 
+                        model.saved_weights[c] = (model.saved_weights[c] *
                                                   wpast_j
                                                   + new_w) / (wpast_j + 1)
                     else:
@@ -505,29 +514,9 @@ def set_task_layer(self, strategy, step_info: IStepInfo):
         :param step_info: the step info object.
         :return: None
         """
-        # TODO: better management of testing flow (especially head expansion)
-        # Main idea for the testing flow: just discard (not store) the layer for
-        # tasks that were not encountered during a training flow.
-        # Can use existing facilities to know if current flow is test or train!
-        # Also, do not expand layers during testing !OR! create a new expanded
-        # layer (which get discarded) by setting the weights of not already
-        # encountered classes to 0!
-        # --- end of dev comment ---
-
-        # compute max label
-        step_id = step_info.current_step
-        train_step = step_info.scenario.train_stream[step_id]
-        # Scenarios may only contain a single complete test set
-        if len(step_info.scenario.test_stream) > 1:
-            test_step = step_info.scenario.test_stream[step_id]
-        else:
-            test_step = step_info.scenario.test_stream[0]
-        train_dataset = train_step.dataset
-        task_label = train_step.task_label
-        test_dataset = test_step.dataset
 
-        n_output_units = max(max(train_dataset.targets),
-                             max(test_dataset.targets)) + 1
+        task_label = step_info.task_label
+        n_output_units = max(step_info.dataset.targets) + 1
 
         if task_label not in self.task_layers:
             # create head for unseen tasks
diff --git a/examples/replay_plugin.py b/examples/replay_plugin.py
--- a/examples/replay_plugin.py
+++ b/examples/replay_plugin.py
@@ -26,14 +26,13 @@
 from torchvision import transforms
 from torchvision.datasets import MNIST
 from torchvision.transforms import ToTensor, RandomCrop
-
-from avalanche.benchmarks.scenarios import DatasetPart, \
-    create_nc_single_dataset_sit_scenario, NCBatchInfo
+import torch.optim.lr_scheduler
+from avalanche.benchmarks import nc_scenario
 from avalanche.evaluation import EvalProtocol
 from avalanche.evaluation.metrics import ACC, CF, RAMU, CM
 from avalanche.extras.models import SimpleMLP
 from avalanche.training.strategies import Naive
-from avalanche.training.plugins import ReplayPlugin
+from avalanche.training.plugins import ReplayPlugin, MultiHeadPlugin
 
 
 def main():
@@ -59,41 +58,39 @@ def main():
                         download=True, transform=train_transform)
     mnist_test = MNIST('./data/mnist', train=False,
                        download=True, transform=test_transform)
-    nc_scenario = create_nc_single_dataset_sit_scenario(
-        mnist_train, mnist_test, n_batches, shuffle=True, seed=1234)
+    scenario = nc_scenario(
+        mnist_train, mnist_test, 5, task_labels=False, seed=1234)
     # ---------
 
     # MODEL CREATION
-    model = SimpleMLP(num_classes=nc_scenario.n_classes)
+    model = SimpleMLP(num_classes=scenario.n_classes)
 
     # DEFINE THE EVALUATION PROTOCOL
     evaluation_protocol = EvalProtocol(
-        metrics=[ACC(num_class=nc_scenario.n_classes),  # Accuracy metric
-                 CF(num_class=nc_scenario.n_classes),  # Catastrophic forgetting
+        metrics=[ACC(num_class=scenario.n_classes),  # Accuracy metric
+                 CF(num_class=scenario.n_classes),  # Catastrophic forgetting
                  RAMU(),  # Ram usage
                  CM()],  # Confusion matrix
         tb_logdir='../logs/mnist_test_sit')
 
     # CREATE THE STRATEGY INSTANCE (NAIVE)
-    cl_strategy = Naive(model, CrossEntropyLoss(), SGD(model.parameters(), lr=0.001, momentum=0.9),
+    cl_strategy = Naive(model, torch.optim.Adam(model.parameters(), lr=0.001),
+        CrossEntropyLoss(),
         train_mb_size=100, train_epochs=4, test_mb_size=100,
         evaluation_protocol=evaluation_protocol, device=device,
-        plugins=[ReplayPlugin(mem_size=10000)]
+        plugins=[ReplayPlugin(mem_size=10000), MultiHeadPlugin(model)]
     )
 
     # TRAINING LOOP
     print('Starting experiment...')
     results = []
-    batch_info: NCBatchInfo
-    for batch_info in nc_scenario:
-        # print("Start of step ", cl_strategy.training_step_id)
-
+    for batch_info in scenario.train_stream:
+        print("Start of step ", batch_info.current_step)
         cl_strategy.train(batch_info, num_workers=4)
         print('Training completed')
 
         print('Computing accuracy on the whole test set')
-        results.append(cl_strategy.test(batch_info, DatasetPart.COMPLETE,
-                                        num_workers=4))
+        results.append(cl_strategy.test(scenario.test_stream, num_workers=4))
 
 
 if __name__ == '__main__':


Test patch: