Repository: continualai/avalanche
Issue: concat_datasets_sequentially discards classes

When using `concat_datasets_sequentially` the resulting datasets have less classes than the union of the classes in the input dataset

```python
import torch
from torch.utils.data import TensorDataset
from avalanche.benchmarks.utils import concat_datasets_sequentially, AvalancheDataset

#  4 training datasets, 4 test datasets with targets [0,1], [2,3], [4,5], [0,1] respectively. 
#  The final output should remap classes in [0-7]

train = [AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(0, 2, (20,)))), 
            AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(2, 4, (20,)))), 
            AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(4, 6, (20,)))), 
            AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(0, 2, (20,))))]

test = [AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(0, 2, (20,)))), 
            AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(2, 4, (20,)))), 
            AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(4, 6, (20,)))), 
            AvalancheDataset(TensorDataset(torch.randn(20,10), torch.randint(0, 2, (20,))))]

final_train, final_test, classes = concat_datasets_sequentially(train, test)

# this is ok -> [[0, 1], [2, 3], [4, 5], [6, 7]]
print(classes) 

# THIS IS WRONG -> {0, 1}, expected {0, 1, 2, 3, 4, 5, 6, 7}
print(set([el[1] for el in final_train]))  
```


Changed files:
avalanche/benchmarks/utils/avalanche_dataset.py, tests/test_avalanche_dataset.py

diff --git a/avalanche/benchmarks/utils/avalanche_dataset.py b/avalanche/benchmarks/utils/avalanche_dataset.py
index fa27b479..f2e9ef2a 100644
--- a/avalanche/benchmarks/utils/avalanche_dataset.py
+++ b/avalanche/benchmarks/utils/avalanche_dataset.py
@@ -2004,20 +2004,34 @@ def concat_datasets_sequentially(
 
     new_class_ids_per_dataset = []
     for dataset_idx in range(len(train_dataset_list)):
+    
+        # Get the train and test sets of the dataset
+        train_set = train_dataset_list[dataset_idx]
+        test_set = test_dataset_list[dataset_idx]
+        
+        # Get the classes in the dataset
+        dataset_classes = set(map(int, train_set.targets))
+        
         # The class IDs for this dataset will be in range
         # [n_classes_in_previous_datasets,
         #       n_classes_in_previous_datasets + classes_in_this_dataset)
-        class_mapping = list(
+        new_classes = list(
             range(next_remapped_idx,
                   next_remapped_idx + classes_per_dataset[dataset_idx]))
-        new_class_ids_per_dataset.append(class_mapping)
-
-        train_set = train_dataset_list[dataset_idx]
-        test_set = test_dataset_list[dataset_idx]
-
+        new_class_ids_per_dataset.append(new_classes)
+        
         # AvalancheSubset is used to apply the class IDs transformation.
         # Remember, the class_mapping parameter must be a list in which:
         # new_class_id = class_mapping[original_class_id]
+        # Hence, a list of size equal to the maximum class index is created
+        # Only elements corresponding to the present classes are remapped
+        class_mapping = [-1] * (max(dataset_classes) + 1)
+        j = 0
+        for i in dataset_classes:
+            class_mapping[i] = new_classes[j]
+            j += 1        
+
+        # Create remapped datasets and append them to the final list
         remapped_train_datasets.append(
             AvalancheSubset(train_set, class_mapping=class_mapping))
         remapped_test_datasets.append(
diff --git a/tests/test_avalanche_dataset.py b/tests/test_avalanche_dataset.py
index 81f4c736..f5e0f2da 100644
--- a/tests/test_avalanche_dataset.py
+++ b/tests/test_avalanche_dataset.py
@@ -23,7 +23,7 @@ from avalanche.benchmarks.scenarios.generic_benchmark_creation import \\
     create_generic_benchmark_from_tensor_lists
 from avalanche.benchmarks.utils import AvalancheDataset, \\
     AvalancheSubset, AvalancheConcatDataset, AvalancheDatasetType, \\
-    AvalancheTensorDataset
+    AvalancheTensorDataset, concat_datasets_sequentially
 from avalanche.benchmarks.utils.dataset_utils import ConstantSequence
 from avalanche.training.utils import load_all_dataset
 import random
@@ -1325,6 +1325,41 @@ class AvalancheDatasetTests(unittest.TestCase):
         self.assertTrue(torch.equal(tensor_t,
                                     leaf[d_sz*dataset_hierarchy_depth:][2]))
 
+    def test_avalanche_concat_datasets_sequentially(self):
+        # create list of training datasets
+        train = [AvalancheDataset(TensorDataset(torch.randn(20, 10),
+                                                torch.randint(0, 2, (20,)))),
+                 AvalancheDataset(TensorDataset(torch.randn(20, 10),
+                                                torch.randint(2, 4, (20,)))),
+                 AvalancheDataset(TensorDataset(torch.randn(20, 10),
+                                                torch.randint(4, 6, (20,)))),
+                 AvalancheDataset(TensorDataset(torch.randn(20, 10),
+                                                torch.randint(0, 2, (20,))))]
+
+        # create list of test datasets
+        test = [AvalancheDataset(TensorDataset(torch.randn(20, 10), 
+                                               torch.randint(0, 2, (20,)))),
+                AvalancheDataset(TensorDataset(torch.randn(20, 10), 
+                                               torch.randint(2, 4, (20,)))),
+                AvalancheDataset(TensorDataset(torch.randn(20, 10), 
+                                               torch.randint(4, 6, (20,)))),
+                AvalancheDataset(TensorDataset(torch.randn(20, 10), 
+                                               torch.randint(0, 2, (20,))))]
+
+        # concatenate datasets
+        final_train, _, classes = concat_datasets_sequentially(train, test)
+
+        # merge all classes into a single list
+        classes_all = []
+        for class_list in classes:
+            classes_all.extend(class_list)
+
+        # get the target set of classes 
+        target_classes = list(set(map(int, final_train.targets)))
+
+        # test for correctness
+        self.assertEqual(classes_all, target_classes)
+
 
 class TransformationSubsetTests(unittest.TestCase):
     def test_avalanche_subset_transform(self):