Fix issue ContinualAI#774

lrzpellegrini · Jul 19, 2023 · 960d013 · 960d013
1 parent 435b40d
commit 960d013
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/avalanche/benchmarks/classic/core50.py b/avalanche/benchmarks/classic/core50.py
@@ -156,6 +156,20 @@ def CORe50(
         eval_transform=eval_transform,
     )
 
+    if scenario == "nc":
+        n_classes_per_exp = []
+        classes_order = []
+        for exp in benchmark_obj.train_stream:
+            exp_dataset = exp.dataset
+            unique_targets = list(
+                sorted(set(int(x) for x in exp_dataset.targets))  # type: ignore
+            )
+            n_classes_per_exp.append(len(unique_targets))
+            classes_order.extend(unique_targets)
+        setattr(benchmark_obj, "n_classes_per_exp", n_classes_per_exp)
+        setattr(benchmark_obj, "classes_order", classes_order)
+    setattr(benchmark_obj, "n_classes", 50 if object_lvl else 10)
+
     return benchmark_obj
 
 

diff --git a/tests/test_core50.py b/tests/test_core50.py
@@ -38,6 +38,11 @@ def test_core50_nc_benchmark(self):
         classes_in_test = benchmark_instance.classes_in_experience["test"][0]
         self.assertSetEqual(set(range(50)), set(classes_in_test))
 
+        # Regression tests for issue #774
+        self.assertSequenceEqual([10] + ([5] * 8), benchmark_instance.n_classes_per_exp)
+        self.assertSetEqual(set(range(50)), set(benchmark_instance.classes_order))
+        self.assertEqual(50, len(benchmark_instance.classes_order))
+
 
 if __name__ == "__main__":
     unittest.main()