[Feature/datasets] Improve dataset management (#33)

* new implementation of datasets

* update dataset implementation

* update dataset implementation

* update dataset implementation and remove reload

* update dataset implementation (all datasets)

* include dataset mode

* include dataset mode

* fix dataset parameters

* optimize dataset management

* clean dataset resources

* improve nebuladataset

Author: enriquetomasmb

docs/_prebuilt/developerguide.md

@@ -259,7 +259,6 @@ First, you must add the Dataset option in the frontend. Adding the Dataset optio
             "EMNIST": ["MLP", "CNN"],
             "CIFAR10": ["CNN", "CNNv2", "CNNv3", "ResNet9", "fastermobilenet", "simplemobilenet"],
             "CIFAR100": ["CNN"],
-            "KITSUN": ["MLP"],
         }
         var datasetSelect = document.getElementById("datasetSelect");
         var modelSelect = document.getElementById("modelSelect");

nebula/addons/trustworthiness/factsheet.py

@@ -19,8 +19,8 @@
     get_feature_importance_cv,
 )
 from nebula.addons.trustworthiness.utils import check_field_filled, count_class_samples, get_entropy, read_csv
-from nebula.core.models.mnist.cnn import CIFAR10ModelCNN, CIFAR10TorchModelCNN, MNISTModelCNN, MNISTTorchModelCNN
-from nebula.core.models.mnist.mlp import MNISTModelMLP, MNISTTorchModelMLP, SyscallModelMLP, SyscallTorchModelMLP
+from nebula.core.models.mnist.cnn import MNISTModelCNN
+from nebula.core.models.mnist.mlp import MNISTModelMLP
 
 dirname = os.path.dirname(__file__)
 
@@ -119,10 +119,8 @@ def populate_factsheet_pre_train(self, data, scenario_name):
                         model = MNISTModelMLP()
                     elif dataset == "MNIST" and algorithm == "CNN":
                         model = MNISTModelCNN()
-                    elif dataset == "Syscall" and algorithm == "MLP":
-                        model = SyscallModelMLP()
                     else:
-                        model = CIFAR10ModelCNN()
+                        model = MNISTModelCNN()
 
                     factsheet["configuration"]["learning_rate"] = model.get_learning_rate()
                     factsheet["configuration"]["trainable_param_num"] = model.count_parameters()
@@ -225,8 +223,6 @@ def populate_factsheet_post_train(self, scenario):
                     pytorch_model = MNISTTorchModelMLP()
                 elif dataset == "MNIST" and model == "CNN":
                     pytorch_model = MNISTTorchModelCNN()
-                elif dataset == "Syscall" and model == "MLP":
-                    pytorch_model = SyscallTorchModelMLP()
                 else:
                     pytorch_model = CIFAR10TorchModelCNN()
 

nebula/controller.py

@@ -193,18 +193,20 @@ def verify_nodes_ports(self, src_path):
         try:
             port_mapping = {}
             new_port_start = 50000
-            
+
             participant_files = sorted(
                 f for f in os.listdir(scenario_path) if f.endswith(".json") and f.startswith("participant")
             )
-            
+
             for filename in participant_files:
                 file_path = os.path.join(scenario_path, filename)
                 with open(file_path) as json_file:
                     node = json.load(json_file)
                 current_port = node["network_args"]["port"]
                 port_mapping[current_port] = SocketUtils.find_free_port(start_port=new_port_start)
-                logging.info(f"Participant file: {filename} | Current port: {current_port} | New port: {port_mapping[current_port]}")
+                logging.info(
+                    f"Participant file: {filename} | Current port: {current_port} | New port: {port_mapping[current_port]}"
+                )
                 new_port_start = port_mapping[current_port] + 1
 
             for filename in participant_files:

nebula/core/datasets/cifar10/cifar10.py

@@ -1,36 +1,63 @@
 import os
 
+from PIL import Image
 from torchvision import transforms
 from torchvision.datasets import CIFAR10
 
-from nebula.core.datasets.nebuladataset import NebulaDataset
+from nebula.core.datasets.nebuladataset import NebulaDataset, NebulaPartitionHandler
+
+
+class CIFAR10PartitionHandler(NebulaPartitionHandler):
+    def __init__(self, file_path, prefix, mode):
+        super().__init__(file_path, prefix, mode)
+
+        # Custom transform for CIFAR10
+        mean = (0.4914, 0.4822, 0.4465)
+        std = (0.2471, 0.2435, 0.2616)
+        self.transform = transforms.Compose([
+            transforms.RandomCrop(32, padding=4),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize(mean, std, inplace=True),
+        ])
+
+    def __getitem__(self, idx):
+        img, target = super().__getitem__(idx)
+
+        img = Image.fromarray(img)
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
 
 
 class CIFAR10Dataset(NebulaDataset):
     def __init__(
         self,
         num_classes=10,
-        partition_id=0,
         partitions_number=1,
         batch_size=32,
         num_workers=4,
         iid=True,
         partition="dirichlet",
         partition_parameter=0.5,
         seed=42,
-        config=None,
+        config_dir=None,
     ):
         super().__init__(
             num_classes=num_classes,
-            partition_id=partition_id,
             partitions_number=partitions_number,
             batch_size=batch_size,
             num_workers=num_workers,
             iid=iid,
             partition=partition,
             partition_parameter=partition_parameter,
             seed=seed,
-            config=config,
+            config_dir=config_dir,
         )
 
     def initialize_dataset(self):
@@ -40,39 +67,15 @@ def initialize_dataset(self):
         if self.test_set is None:
             self.test_set = self.load_cifar10_dataset(train=False)
 
-        # All nodes have the same test set (indices are the same for all nodes)
-        self.test_indices_map = list(range(len(self.test_set)))
-
-        # Depending on the iid flag, generate a non-iid or iid map of the train set
-        if self.iid:
-            self.train_indices_map = self.generate_iid_map(self.train_set, self.partition, self.partition_parameter)
-            self.local_test_indices_map = self.generate_iid_map(self.test_set, self.partition, self.partition_parameter)
-        else:
-            self.train_indices_map = self.generate_non_iid_map(self.train_set, self.partition, self.partition_parameter)
-            self.local_test_indices_map = self.generate_non_iid_map(
-                self.test_set, self.partition, self.partition_parameter
-            )
-
-        print(f"Length of train indices map: {len(self.train_indices_map)}")
-        print(f"Lenght of test indices map (global): {len(self.test_indices_map)}")
-        print(f"Length of test indices map (local): {len(self.local_test_indices_map)}")
+        self.data_partitioning(plot=True)
 
     def load_cifar10_dataset(self, train=True):
-        mean = (0.4914, 0.4822, 0.4465)
-        std = (0.2471, 0.2435, 0.2616)
-        apply_transforms = transforms.Compose([
-            transforms.RandomCrop(32, padding=4),
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean, std, inplace=True),
-        ])
         data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
         os.makedirs(data_dir, exist_ok=True)
         return CIFAR10(
             data_dir,
             train=train,
             download=True,
-            transform=apply_transforms,
         )
 
     def generate_non_iid_map(self, dataset, partition="dirichlet", partition_parameter=0.5):
@@ -83,11 +86,7 @@ def generate_non_iid_map(self, dataset, partition="dirichlet", partition_paramet
         else:
             raise ValueError(f"Partition {partition} is not supported for Non-IID map")
 
-        if self.partition_id == 0:
-            self.plot_data_distribution(dataset, partitions_map)
-            self.plot_all_data_distribution(dataset, partitions_map)
-
-        return partitions_map[self.partition_id]
+        return partitions_map
 
     def generate_iid_map(self, dataset, partition="balancediid", partition_parameter=2):
         if partition == "balancediid":
@@ -97,8 +96,4 @@ def generate_iid_map(self, dataset, partition="balancediid", partition_parameter
         else:
             raise ValueError(f"Partition {partition} is not supported for IID map")
 
-        if self.partition_id == 0:
-            self.plot_data_distribution(dataset, partitions_map)
-            self.plot_all_data_distribution(dataset, partitions_map)
-
-        return partitions_map[self.partition_id]
+        return partitions_map

nebula/core/datasets/cifar100/cifar100.py

@@ -1,36 +1,63 @@
 import os
 
+from PIL import Image
 from torchvision import transforms
 from torchvision.datasets import CIFAR100
 
-from nebula.core.datasets.nebuladataset import NebulaDataset
+from nebula.core.datasets.nebuladataset import NebulaDataset, NebulaPartitionHandler
+
+
+class CIFAR100PartitionHandler(NebulaPartitionHandler):
+    def __init__(self, file_path, prefix, mode):
+        super().__init__(file_path, prefix, mode)
+
+        # Custom transform for CIFAR100
+        mean = (0.4914, 0.4822, 0.4465)
+        std = (0.2471, 0.2435, 0.2616)
+        self.transform = transforms.Compose([
+            transforms.RandomCrop(32, padding=4),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize(mean, std, inplace=True),
+        ])
+
+    def __getitem__(self, idx):
+        img, target = super().__getitem__(idx)
+
+        img = Image.fromarray(img)
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
 
 
 class CIFAR100Dataset(NebulaDataset):
     def __init__(
         self,
         num_classes=100,
-        partition_id=0,
         partitions_number=1,
         batch_size=32,
         num_workers=4,
         iid=True,
         partition="dirichlet",
         partition_parameter=0.5,
         seed=42,
-        config=None,
+        config_dir=None,
     ):
         super().__init__(
             num_classes=num_classes,
-            partition_id=partition_id,
             partitions_number=partitions_number,
             batch_size=batch_size,
             num_workers=num_workers,
             iid=iid,
             partition=partition,
             partition_parameter=partition_parameter,
             seed=seed,
-            config=config,
+            config_dir=config_dir,
         )
 
     def initialize_dataset(self):
@@ -40,37 +67,13 @@ def initialize_dataset(self):
         if self.test_set is None:
             self.test_set = self.load_cifar100_dataset(train=False)
 
-        # All nodes have the same test set (indices are the same for all nodes)
-        self.test_indices_map = list(range(len(self.test_set)))
-
-        # Depending on the iid flag, generate a non-iid or iid map of the train set
-        if self.iid:
-            self.train_indices_map = self.generate_iid_map(self.train_set, self.partition, self.partition_parameter)
-            self.local_test_indices_map = self.generate_iid_map(self.test_set, self.partition, self.partition_parameter)
-        else:
-            self.train_indices_map = self.generate_non_iid_map(self.train_set, self.partition, self.partition_parameter)
-            self.local_test_indices_map = self.generate_non_iid_map(
-                self.test_set, self.partition, self.partition_parameter
-            )
-
-        print(f"Length of train indices map: {len(self.train_indices_map)}")
-        print(f"Lenght of test indices map (global): {len(self.test_indices_map)}")
-        print(f"Length of test indices map (local): {len(self.local_test_indices_map)}")
+        self.data_partitioning(plot=True)
 
     def load_cifar100_dataset(self, train=True):
-        mean = (0.4914, 0.4822, 0.4465)
-        std = (0.2471, 0.2435, 0.2616)
-        apply_transforms = transforms.Compose([
-            transforms.RandomCrop(32, padding=4),
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean, std, inplace=True),
-        ])
         return CIFAR100(
             os.path.join(os.path.dirname(os.path.abspath(__file__)), "data"),
             train=train,
             download=True,
-            transform=apply_transforms,
         )
 
     def generate_non_iid_map(self, dataset, partition="dirichlet", partition_parameter=0.5):
@@ -81,11 +84,7 @@ def generate_non_iid_map(self, dataset, partition="dirichlet", partition_paramet
         else:
             raise ValueError(f"Partition {partition} is not supported for Non-IID map")
 
-        if self.partition_id == 0:
-            self.plot_data_distribution(dataset, partitions_map)
-            self.plot_all_data_distribution(dataset, partitions_map)
-
-        return partitions_map[self.partition_id]
+        return partitions_map
 
     def generate_iid_map(self, dataset, partition="balancediid", partition_parameter=2):
         if partition == "balancediid":
@@ -95,8 +94,4 @@ def generate_iid_map(self, dataset, partition="balancediid", partition_parameter
         else:
             raise ValueError(f"Partition {partition} is not supported for IID map")
 
-        if self.partition_id == 0:
-            self.plot_data_distribution(dataset, partitions_map)
-            self.plot_all_data_distribution(dataset, partitions_map)
-
-        return partitions_map[self.partition_id]
+        return partitions_map

nebula/core/datasets/datamodule.py

@@ -18,8 +18,6 @@ def __init__(
         test_set,
         test_set_indices,
         local_test_set_indices,
-        partition_id=0,
-        partitions_number=1,
         batch_size=32,
         num_workers=0,
         val_percent=0.1,
@@ -31,8 +29,6 @@ def __init__(
         self.test_set = test_set
         self.test_set_indices = test_set_indices
         self.local_test_set_indices = local_test_set_indices
-        self.partition_id = partition_id
-        self.partitions_number = partitions_number
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.val_percent = val_percent
@@ -79,9 +75,6 @@ def setup(self, stage=None):
             self.global_te_subset = ChangeableSubset(self.test_set, self.test_set_indices)
             self.local_te_subset = ChangeableSubset(self.test_set, self.local_test_set_indices)
 
-            if len(self.test_set) < self.partitions_number:
-                raise ValueError("Too many partitions for the size of the test set.")
-
     def teardown(self, stage=None):
         # Teardown the datasets
         if stage in (None, "fit"):