fix existing bugs

Author: WwZzz

flgo/algorithm/afl.py

@@ -47,7 +47,7 @@ def project(self, p):
             res.append(max(p[i] + lmbd, 0))
         return res
 
-    def global_test(self, dataflag='valid'):
+    def global_test(self, flag='valid'):
         """
         Validate accuracies and losses on clients' local datasets
         :param
@@ -57,7 +57,7 @@ def global_test(self, dataflag='valid'):
         """
         all_metrics = collections.defaultdict(list)
         for c in self.clients:
-            client_metrics = c.test(self.result_model, dataflag)
+            client_metrics = c.test(self.result_model, flag)
             for met_name, met_val in client_metrics.items():
                 all_metrics[met_name].append(met_val)
         return all_metrics

flgo/algorithm/fedbase.py

@@ -318,7 +318,7 @@ def aggregate(self, models: list, *args, **kwargs):
             p = [pk/sump for pk in p]
             return fmodule._model_sum([model_k * pk for model_k, pk in zip(models, p)])
 
-    def global_test(self, dataflag='valid'):
+    def global_test(self, flag='valid'):
         """
         Validate accuracies and losses on clients' local datasets
         :param
@@ -328,7 +328,7 @@ def global_test(self, dataflag='valid'):
         """
         all_metrics = collections.defaultdict(list)
         for c in self.clients:
-            client_metrics = c.test(self.model, dataflag)
+            client_metrics = c.test(self.model, flag)
             for met_name, met_val in client_metrics.items():
                 all_metrics[met_name].append(met_val)
         return all_metrics
@@ -345,7 +345,7 @@ def test(self, model=None, flag='test'):
         data = self.test_data if flag=='test' else self.valid_data
         if data is None: return {}
         else:
-            return self.calculator.test(model, self.test_data, batch_size = self.option['test_batch_size'])
+            return self.calculator.test(model, data, batch_size = self.option['test_batch_size'], num_workers = self.option['num_workers'], pin_memory = self.option['pin_memory'])
 
     def init_algo_para(self, algo_para: dict):
         """
@@ -442,6 +442,7 @@ def __init__(self, option={}):
         # server
         self.server = None
         # actions of different message type
+        self.option = option
         self.actions = {0: self.reply}
 
     def initialize(self):
@@ -481,7 +482,7 @@ def test(self, model, dataflag='valid'):
         """
         dataset = self.train_data if dataflag=='train' else self.valid_data
         if dataset is not None:
-            return self.calculator.test(model, dataset, self.test_batch_size)
+            return self.calculator.test(model, dataset, self.test_batch_size, self.option['num_workers'])
         else:
             return {}
 
@@ -610,7 +611,7 @@ def get_batch_data(self):
         try:
             batch_data = next(self.data_loader)
         except Exception as e:
-            self.data_loader = iter(self.calculator.get_dataloader(self.train_data, batch_size=self.batch_size, num_workers=self.loader_num_workers))
+            self.data_loader = iter(self.calculator.get_dataloader(self.train_data, batch_size=self.batch_size, num_workers=self.loader_num_workers, pin_memory=self.option['pin_memory']))
             batch_data = next(self.data_loader)
         # clear local DataLoader when finishing local training
         self.current_steps = (self.current_steps+1) % self.num_steps

flgo/benchmark/toolkits/cv/horizontal/image_classification.py

@@ -82,37 +82,36 @@ def compute_loss(self, model, data):
         return {'loss': loss}
 
     @torch.no_grad()
-    def test(self, model, dataset, batch_size=64, num_workers=0):
+    def test(self, model, dataset, batch_size=64, num_workers=0, pin_memory=False):
         """
         Metric = [mean_accuracy, mean_loss]
         :param model:
         :param dataset:
         :param batch_size:
         :return: [mean_accuracy, mean_loss]
         """
-        with torch.no_grad():
-            model.eval()
-            if batch_size==-1:batch_size=len(dataset)
-            data_loader = self.get_dataloader(dataset, batch_size=batch_size, num_workers=num_workers)
-            total_loss = 0.0
-            num_correct = 0
-            for batch_id, batch_data in enumerate(data_loader):
-                batch_data = self.to_device(batch_data)
-                outputs = model(batch_data[0])
-                batch_mean_loss = self.criterion(outputs, batch_data[-1]).item()
-                y_pred = outputs.data.max(1, keepdim=True)[1]
-                correct = y_pred.eq(batch_data[-1].data.view_as(y_pred)).long().cpu().sum()
-                num_correct += correct.item()
-                total_loss += batch_mean_loss * len(batch_data[-1])
+        model.eval()
+        if batch_size==-1:batch_size=len(dataset)
+        data_loader = self.get_dataloader(dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory)
+        total_loss = 0.0
+        num_correct = 0
+        for batch_id, batch_data in enumerate(data_loader):
+            batch_data = self.to_device(batch_data)
+            outputs = model(batch_data[0])
+            batch_mean_loss = self.criterion(outputs, batch_data[-1]).item()
+            y_pred = outputs.data.max(1, keepdim=True)[1]
+            correct = y_pred.eq(batch_data[-1].data.view_as(y_pred)).long().cpu().sum()
+            num_correct += correct.item()
+            total_loss += batch_mean_loss * len(batch_data[-1])
         return {'accuracy': 1.0*num_correct/len(dataset), 'loss':total_loss/len(dataset)}
 
     def to_device(self, data):
         return data[0].to(self.device), data[1].to(self.device)
 
-    def get_dataloader(self, dataset, batch_size=64, shuffle=True, num_workers=0):
+    def get_dataloader(self, dataset, batch_size=64, shuffle=True, num_workers=0, pin_memory=False, drop_last=False):
         if self.DataLoader == None:
             raise NotImplementedError("DataLoader Not Found.")
-        return self.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
+        return self.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory, drop_last=drop_last)
 
 class GeneralGenerator(BasicTaskGenerator):
     def __init__(self, benchmark, rawdata_path):

flgo/experiment/analyzer.py

@@ -101,8 +101,14 @@ def set_communication_round(self):
 
     def set_client_id(self):
         with open(os.path.join(self.task, 'info')) as inf:
-            num_clients = json.load(inf)['num_clients']
-        self.data['client_id'] = [cid for cid in range(int(num_clients))]
+            task_info = json.load(inf)
+            if 'num_clients' in task_info.keys():
+                N = int(task_info['num_clients'])
+            elif 'num_parties' in task_info.keys():
+                N = int(task_info['num_parties'])
+            else:
+                N = 0
+        self.data['client_id'] = [cid for cid in range(N)]
 
     def set_legend(self, legend_with = []):
         if len(legend_with)==0: self.data['label'] = []

flgo/utils/fflow.py

@@ -34,7 +34,7 @@
 sample_list=['uniform', 'md', 'full', 'uniform_available', 'md_available', 'full_available']
 agg_list=['uniform', 'weighted_scale', 'weighted_com']
 optimizer_list=['SGD', 'Adam', 'RMSprop', 'Adagrad']
-default_option_dict = {'pretrain': '', 'sample': 'md', 'aggregate': 'uniform', 'num_rounds': 20, 'proportion': 0.2, 'learning_rate_decay': 0.998, 'lr_scheduler': -1, 'early_stop': -1, 'num_epochs': 5, 'num_steps': -1, 'learning_rate': 0.1, 'batch_size': 64.0, 'optimizer': 'SGD', 'momentum': 0, 'weight_decay': 0, 'algo_para': [], 'train_holdout': 0.1, 'test_holdout': 0.0, 'seed': 0, 'gpu': [], 'server_with_cpu': False, 'num_parallels': 1, 'num_workers': 0, 'test_batch_size': 512, 'simulator': 'default_simulator', 'availability': 'IDL', 'connectivity': 'IDL', 'completeness': 'IDL', 'responsiveness': 'IDL', 'logger': 'basic_logger', 'log_level': 'INFO', 'log_file': False, 'no_log_console': False, 'no_overwrite': False, 'eval_interval': 1}
+default_option_dict = {'pretrain': '', 'sample': 'md', 'aggregate': 'uniform', 'num_rounds': 20, 'proportion': 0.2, 'learning_rate_decay': 0.998, 'lr_scheduler': -1, 'early_stop': -1, 'num_epochs': 5, 'num_steps': -1, 'learning_rate': 0.1, 'batch_size': 64.0, 'optimizer': 'SGD', 'momentum': 0, 'weight_decay': 0, 'algo_para': [], 'train_holdout': 0.1, 'test_holdout': 0.0, 'seed': 0, 'gpu': [], 'server_with_cpu': False, 'num_parallels': 1, 'num_workers': 0, 'pin_memory':False,'test_batch_size': 512, 'simulator': 'default_simulator', 'availability': 'IDL', 'connectivity': 'IDL', 'completeness': 'IDL', 'responsiveness': 'IDL', 'logger': 'basic_logger', 'log_level': 'INFO', 'log_file': False, 'no_log_console': False, 'no_overwrite': False, 'eval_interval': 1}
 
 class GlobalVariable:
     """this class is to create a buffer space for sharing variables across different parties for each runner respectively in a single machine"""
@@ -99,6 +99,7 @@ def read_option_from_command():
     parser.add_argument('--server_with_cpu', help='seed for random initialization;', action="store_true", default=False)
     parser.add_argument('--num_parallels', help="the number of parallels in the clients computing session", type=int, default=1)
     parser.add_argument('--num_workers', help='the number of workers of DataLoader', type=int, default=0)
+    parser.add_argument('--pin_memory', help='pin_memory of DataLoader', action="store_true", default=False)
     parser.add_argument('--test_batch_size', help='the batch_size used in testing phase;', type=int, default=512)
 
     """Simulator Options"""