add iid splitter (#394)

federatedscope/core/auxiliaries/splitter_builder.py

@@ -39,6 +39,9 @@ def get_splitter(config):
     elif config.data.splitter == 'rand_chunk':
         from federatedscope.core.splitters.graph import RandChunkSplitter
         splitter = RandChunkSplitter(client_num, **kwargs)
+    elif config.data.splitter == 'iid':
+        from federatedscope.core.splitters.generic import IIDSplitter
+        splitter = IIDSplitter(client_num)
     else:
         logger.warning(f'Splitter {config.data.splitter} not found.')
         splitter = None

federatedscope/core/splitters/generic/__init__.py

@@ -1,3 +1,4 @@
 from federatedscope.core.splitters.generic.lda_splitter import LDASplitter
+from federatedscope.core.splitters.generic.iid_splitter import IIDSplitter
 
-__all__ = ['LDASplitter']
+__all__ = ['LDASplitter', 'IIDSplitter']

federatedscope/core/splitters/generic/iid_splitter.py

@@ -0,0 +1,17 @@
+import numpy as np
+from federatedscope.core.splitters import BaseSplitter
+
+
+class IIDSplitter(BaseSplitter):
+    def __init__(self, client_num):
+        super(IIDSplitter, self).__init__(client_num)
+
+    def __call__(self, dataset, prior=None):
+        dataset = [ds for ds in dataset]
+        np.random.shuffle(dataset)
+        length = len(dataset)
+        prop = [1.0 / self.client_num for _ in range(self.client_num)]
+        prop = (np.cumsum(prop) * length).astype(int)[:-1]
+        data_list = np.split(dataset, prop)
+        data_list = [x.tolist() for x in data_list]
+        return data_list