Make intention fo distributed datasets more clear

JMGaljaard · Sep 4, 2022 · 949803b · 949803b
1 parent 45d1017
commit 949803b
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 22 deletions.
diff --git a/fltk/datasets/cifar10.py b/fltk/datasets/cifar10.py
@@ -7,7 +7,7 @@
 
 class CIFAR10Dataset(Dataset):
     """
-    CIFAR10 Dataset implementation for Federated learning experiments.
+    CIFAR10 Dataset implementation for Distributed learning experiments.
     """
 
     def __init__(self, config, learning_param, rank: int = 0, world_size: int = None):

diff --git a/fltk/datasets/cifar100.py b/fltk/datasets/cifar100.py
@@ -7,7 +7,7 @@
 
 class CIFAR100Dataset(Dataset):
     """
-    CIFAR100 Dataset implementation for Federated learning experiments.
+    CIFAR100 Dataset implementation for Distributed learning experiments.
     """
 
     DEFAULT_TRANSFORM = transforms.Compose([

diff --git a/fltk/datasets/dataset.py b/fltk/datasets/dataset.py
@@ -1,4 +1,6 @@
 from __future__ import annotations
+
+import abc
 from abc import abstractmethod
 
 import torch
@@ -11,7 +13,10 @@
     from fltk.util.config import DistLearningConfig
 
 
-class Dataset:
+class Dataset(abc.ABC):
+    """
+    Dataset implementation for Distributed learning experiments.
+    """
 
     def __init__(self, config, learning_params: DistLearningConfig, rank: int, world_size: int):
         self.config = config
@@ -39,24 +44,6 @@ def get_test_dataset(self):
         """
         return self.test_loader
 
-    @abstractmethod
-    def load_train_dataset(self):
-        """
-        Loads & returns the training dataset.
-
-        :return: tuple
-        """
-        raise NotImplementedError("load_train_dataset() isn't implemented")
-
-    @abstractmethod
-    def load_test_dataset(self):
-        """
-        Loads & returns the test dataset.
-
-        :return: tuple
-        """
-        raise NotImplementedError("load_test_dataset() isn't implemented")
-
     def get_train_loader(self, **kwargs):
         """
         Return the data loader for the train dataset.
@@ -77,7 +64,24 @@ def get_test_loader(self, **kwargs):
         """
         return self.test_loader
 
-    @staticmethod
+    @abstractmethod
+    def load_train_dataset(self):
+        """
+        Loads & returns the training dataset.
+
+        :return: tuple
+        """
+        raise NotImplementedError("load_train_dataset() isn't implemented")
+
+    @abstractmethod
+    def load_test_dataset(self):
+        """
+        Loads & returns the test dataset.
+
+        :return: tuple
+        """
+        raise NotImplementedError("load_test_dataset() isn't implemented")
+
     def get_data_loader_from_data(batch_size, X, Y, **kwargs):
         """
         Get a data loader created from a given set of data.

diff --git a/fltk/datasets/fashion_mnist.py b/fltk/datasets/fashion_mnist.py
@@ -6,6 +6,9 @@
 
 
 class FashionMNISTDataset(Dataset):
+    """
+    FashionMNIST Dataset implementation for Distributed learning experiments.
+    """
 
     def __init__(self, config, learning_param, rank: int = 0, world_size: int = None):
         super(FashionMNISTDataset, self).__init__(config, learning_param, rank, world_size)

diff --git a/fltk/datasets/mnist.py b/fltk/datasets/mnist.py
@@ -7,6 +7,10 @@
 
 
 class MNIST(Dataset):
+    """
+    MNIST Dataset implementation for Distributed learning experiments.
+    """
+
     DEFAULT_TRANSFORM = transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))