add report for all partitinoer

SMILELab-FL · Dec 7, 2023 · e315584 · e315584
1 parent be21ab4
commit e315584
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/fedlab/utils/dataset/functional.py b/fedlab/utils/dataset/functional.py
@@ -556,18 +556,18 @@ def partition_report(targets, data_indices, class_num=None, verbose=True, file=N
         Check ``data_indices`` may look like:
 
         >>> data_indices
-        {0: array([8, 6, 5, 7, 2]),
-         1: array([ 3, 10, 14,  4,  1]),
-         2: array([13,  9, 12, 11,  0])}
+        {0: array([ 4,  1, 14,  8,  5]),
+         1: array([ 0, 13, 12,  3,  2]),
+         2: array([10,  9,  7,  6, 11])}
 
         Now generate partition report for each client and each class:
 
         >>> partition_report(labels, data_indices, class_num=class_num, verbose=True, file=None)
         Class sample statistics:
-        client,class0,class1,class2,class3,Amount
-        Client   0,0.200,0.00,0.200,0.600,5
-        Client   1,0.400,0.200,0.200,0.200,5
-        Client   2,0.00,0.400,0.400,0.200,5
+           cid  class-0  class-1  class-2  class-3  TotalAmount
+        0    0        3        2        0        0            5
+        1    1        1        1        1        2            5
+        2    2        3        1        1        0            5
 
     """
     if not isinstance(targets, np.ndarray):

diff --git a/fedlab/utils/dataset/partition.py b/fedlab/utils/dataset/partition.py
@@ -266,6 +266,7 @@ def __init__(self, targets, num_clients,
         self.client_dict = self._perform_partition()
         # get sample number count for each client
         self.client_sample_count = F.samples_num_count(self.client_dict, self.num_clients)
+        self.stats_report = F.partition_report(targets, self.client_dict, class_num=self.num_classes, verbose=False)
 
     def _perform_partition(self):
         if self.partition == "noniid-#label":