Repository: continualai/avalanche
Issue: Bad class remapping in NCScenario

The datasets returned by `NCScenario` have a bad class (targets) mapping when `class_ids_from_zero_from_first_step` is set to `True`.

For instance, the dataset will contain classes with random IDs. Instead, when `class_ids_from_zero_from_first_step=True`, we expect class IDs to be in range [0, n_classes_step0) for the first step, [n_classes_step0, n_classes_step0+n_classes_step1) for the second step, etc.

I'm think I already have a fix for this but I have to check if this also applies to the `class_ids_from_zero_in_each_step` parameter.

Unit tests don't catch this bug as they only check if the `classes_in_step` field of `NCScenario` is correct (which it is).  A regression unit test will be added ASAP.

Changed files:
tests/test_nc_sit_scenario.py, avalanche/benchmarks/scenarios/new_classes/nc_scenario.py

diff --git a/avalanche/benchmarks/scenarios/new_classes/nc_scenario.py b/avalanche/benchmarks/scenarios/new_classes/nc_scenario.py
index 189b5d1d..2eca8ef3 100644
--- a/avalanche/benchmarks/scenarios/new_classes/nc_scenario.py
+++ b/avalanche/benchmarks/scenarios/new_classes/nc_scenario.py
@@ -134,8 +134,15 @@ class NCScenario(GenericCLScenario[TrainSet, TestSet, 'NCStepInfo'],
         n_original_classes = max(self.classes_order_original_ids) + 1
 
         self.class_mapping: List[int] = []
-        """ class_mapping stores the class mapping so that
-        mapped_class_id = class_mapping[original_class_id]. """
+        """
+        class_mapping stores the class mapping so that 
+        `mapped_class_id = class_mapping[original_class_id]`. 
+        
+        If the scenario is created with an amount of classes which is less than
+        the amount of all classes in the dataset, then class_mapping will 
+        contain some -1 values corresponding to ignored classes. This can
+        happen when passing a fixed class order to the constructor.
+        """
 
         self.n_classes_per_step: List[int] = []
         """ A list that, for each step (identified by its index/ID),
@@ -271,7 +278,10 @@ class NCScenario(GenericCLScenario[TrainSet, TestSet, 'NCStepInfo'],
             # over all steps
             self.classes_order = list(range(0, self.n_classes))
             self.class_mapping = [-1] * n_original_classes
-            for class_id in range(self.n_classes):
+            for class_id in range(n_original_classes):
+                # This check is needed because, when a fixed class order is
+                # used, the user may have defined an amount of classes less than
+                # the overall amount of classes in the dataset.
                 if class_id in self.classes_order_original_ids:
                     self.class_mapping[class_id] = \\
                         self.classes_order_original_ids.index(class_id)
@@ -417,15 +427,7 @@ class NCStepInfo(GenericStepInfo[NCScenario[TrainSet, TestSet],
         :param origin_stream: The stream from which this step was obtained.
         :param current_step: The current step ID, as an integer.
         """
-        super(NCStepInfo, self).__init__(
-            origin_stream, current_step)
-
-    @property
-    def dataset(self):
-        dataset = GenericStepInfo.dataset.fget(self)
-
-        return TransformationSubset(
-            dataset, class_mapping=self.scenario.class_mapping)
+        super(NCStepInfo, self).__init__(origin_stream, current_step)
 
 
 __all__ = ['NCScenario', 'NCStepInfo']
diff --git a/tests/test_nc_sit_scenario.py b/tests/test_nc_sit_scenario.py
index 7ebb1da3..33cc3b12 100644
--- a/tests/test_nc_sit_scenario.py
+++ b/tests/test_nc_sit_scenario.py
@@ -75,9 +75,58 @@ class SITTests(unittest.TestCase):
         for batch_id in range(4):
             self.assertEqual(2, len(my_nc_scenario.classes_in_step[batch_id]))
             all_classes.extend(my_nc_scenario.classes_in_step[batch_id])
-
         self.assertEqual(list(range(8)), all_classes)
 
+        # Regression test for issue #258
+        for i, step_info in enumerate(my_nc_scenario.train_stream):
+            unique_dataset_classes = sorted(set(step_info.dataset.targets))
+            expected_dataset_classes = list(range(2 * i, 2 * (i+1)))
+
+            self.assertListEqual(expected_dataset_classes,
+                                 unique_dataset_classes)
+            self.assertListEqual(
+                sorted(order[2 * i:2 * (i+1)]),
+                sorted(my_nc_scenario.original_classes_in_step[i]))
+        # End regression test for issue #258
+
+    def test_sit_single_dataset_remap_indexes_each_step(self):
+        order = [2, 3, 5, 8, 9, 1, 4, 6]
+        mnist_train = MNIST('./data/mnist', train=True, download=True)
+        mnist_test = MNIST('./data/mnist', train=False, download=True)
+
+        with self.assertRaises(ValueError):
+            # class_ids_from_zero_* are mutually exclusive
+            nc_scenario(
+                mnist_train, mnist_test, 4, task_labels=False,
+                fixed_class_order=order,
+                class_ids_from_zero_from_first_step=True,
+                class_ids_from_zero_in_each_step=True)
+
+        my_nc_scenario = nc_scenario(
+            mnist_train, mnist_test, 4, task_labels=False,
+            fixed_class_order=order,
+            class_ids_from_zero_in_each_step=True)
+
+        self.assertEqual(4, len(my_nc_scenario.classes_in_step))
+
+        all_classes = []
+        for batch_id in range(4):
+            self.assertEqual(2, len(my_nc_scenario.classes_in_step[batch_id]))
+            all_classes.extend(my_nc_scenario.classes_in_step[batch_id])
+        self.assertEqual(8, len(all_classes))
+        self.assertListEqual([0, 1], sorted(set(all_classes)))
+
+        # Regression test for issue #258
+        for i, step_info in enumerate(my_nc_scenario.train_stream):
+            unique_dataset_classes = sorted(set(step_info.dataset.targets))
+            expected_dataset_classes = [0, 1]
+            self.assertListEqual(expected_dataset_classes,
+                                 unique_dataset_classes)
+            self.assertListEqual(
+                sorted(order[2 * i:2 * (i + 1)]),
+                sorted(my_nc_scenario.original_classes_in_step[i]))
+        # End regression test for issue #258
+
     def test_sit_single_dataset_reproducibility_data(self):
         mnist_train = MNIST('./data/mnist', train=True, download=True)
         mnist_test = MNIST('./data/mnist', train=False, download=True)