Distributed support (rework)

avalanche/benchmarks/classic/cmnist.py

@@ -29,6 +29,7 @@
 )
 from avalanche.benchmarks.datasets import default_dataset_location
 from avalanche.benchmarks.utils import AvalancheDataset
+from avalanche.distributed import DistributedHelper
 
 _default_mnist_train_transform = Compose(
     [ToTensor(), Normalize((0.1307,), (0.3081,))]
@@ -394,9 +395,12 @@ def _get_mnist_dataset(dataset_root):
     if dataset_root is None:
         dataset_root = default_dataset_location("mnist")
 
-    train_set = MNIST(root=dataset_root, train=True, download=True)
+    with DistributedHelper.main_process_first():
+        train_set = MNIST(root=dataset_root,
+                          train=True, download=True)
 
-    test_set = MNIST(root=dataset_root, train=False, download=True)
+        test_set = MNIST(root=dataset_root,
+                         train=False, download=True)
 
     return train_set, test_set
 

avalanche/benchmarks/utils/collate_functions.py

@@ -0,0 +1,69 @@
+################################################################################
+# Copyright (c) 2021 ContinualAI.                                              #
+# Copyrights licensed under the MIT License.                                   #
+# See the accompanying LICENSE file for terms.                                 #
+#                                                                              #
+# Date: 21-04-2022                                                             #
+# Author(s): Antonio Carta, Lorenzo Pellegrini                                 #
+# E-mail: [email protected]                                              #
+# Website: avalanche.continualai.org                                           #
+################################################################################
+
+import itertools
+from collections import defaultdict
+
+import torch
+
+
+def classification_collate_mbatches_fn(mbatches):
+    """Combines multiple mini-batches together.
+
+        Concatenates each tensor in the mini-batches along dimension 0 (usually
+        this is the batch size).
+
+        :param mbatches: sequence of mini-batches.
+        :return: a single mini-batch
+        """
+    batch = []
+    for i in range(len(mbatches[0])):
+        t = classification_single_values_collate_fn(
+            [el[i] for el in mbatches], i)
+        batch.append(t)
+    return batch
+
+
+def classification_single_values_collate_fn(values_list, index):
+    return torch.cat(values_list, dim=0)
+
+
+def detection_collate_fn(batch):
+    """
+    Collate function used when loading detection datasets using a DataLoader.
+    """
+    return tuple(zip(*batch))
+
+
+def detection_collate_mbatches_fn(mbatches):
+    """
+    Collate function used when loading detection datasets using a DataLoader.
+    """
+    lists_dict = defaultdict(list)
+    for mb in mbatches:
+        for mb_elem_idx, mb_elem in enumerate(mb):
+            lists_dict[mb_elem_idx].append(mb_elem)
+
+    lists = []
+    for mb_elem_idx in range(max(lists_dict.keys()) + 1):
+        lists.append(list(itertools.chain.from_iterable(
+            lists_dict[mb_elem_idx]
+        )))
+
+    return lists
+
+
+__all__ = [
+    'classification_collate_mbatches_fn',
+    'classification_single_values_collate_fn',
+    'detection_collate_fn',
+    'detection_collate_mbatches_fn'
+]