Update Homo LR

Signed-off-by: cwj <talkingwallace@sohu.com> Signed-off-by: weiwee <wbwmat@gmail.com>
FederatedAI · Jul 21, 2023 · 7bd3c59 · 7bd3c59
1 parent 0a3b93c
commit 7bd3c59
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 20 deletions.
diff --git a/python/fate/ml/glm/homo/lr/client.py b/python/fate/ml/glm/homo/lr/client.py
@@ -424,7 +424,7 @@ def fit(self, ctx: Context, train_data: DataFrame,
             self.trainer.set_local_mode()
         self.trainer.train()
 
-        logger.info('training finished')
+        logger.info('homo lr fit done')
 
     def predict(self, ctx: Context, predict_data: DataFrame) -> DataFrame:
 

diff --git a/python/fate/ml/glm/homo/lr/test/test_fed_lr.py b/python/fate/ml/glm/homo/lr/test/test_fed_lr.py
@@ -50,11 +50,17 @@ def create_ctx(local):
             dtype="object")
 
         data = reader.to_frame(ctx, df)
+        client = HomoLRClient(
+            50, 800, optimizer_param={
+                'method': 'adam', 'penalty': 'l1', 'aplha': 0.1, 'optimizer_para': {
+                    'lr': 0.1}}, init_param={
+                        'method': 'random', 'fill_val': 1.0})
 
+        client.fit(ctx, data)
 
     elif sys.argv[1] == "host":
 
-        ctx = create_ctx(guest)
+        ctx = create_ctx(host)
         df = pd.read_csv(
         '../../../../../../../examples/data/breast_homo_host.csv')
         df['sample_id'] = [i for i in range(len(df))]
@@ -66,6 +72,15 @@ def create_ctx(local):
             dtype="object")
 
         data = reader.to_frame(ctx, df)
-
+        client = HomoLRClient(
+            50, 800, optimizer_param={
+                'method': 'adam', 'penalty': 'l1', 'aplha': 0.1, 'optimizer_para': {
+                    'lr': 0.1}}, init_param={
+                        'method': 'random', 'fill_val': 1.0})
+
+        client.fit(ctx, data)
     else:
-        ctx = create_ctx(arbiter)
+
+        ctx = create_ctx(arbiter)
+        server = HomoLRServer()
+        server.fit(ctx)
diff --git a/python/fate/ml/nn/algo/homo/fedavg.py b/python/fate/ml/nn/algo/homo/fedavg.py
@@ -24,19 +24,7 @@
 
 @dataclass
 class FedAVGArguments(FedArguments):
-
-    """
-    The arguemnt for FedAVG algorithm, used in FedAVGClient and FedAVGServer.
-
-    Attributes:
-        weighted_aggregate: bool
-            Whether to use weighted aggregation or not.
-        secure_aggregate: bool
-            Whether to use secure aggregation or not.
-    """
-
-    weighted_aggregate: bool = field(default=True)
-    secure_aggregate: bool = field(default=False)
+    pass
 
 
 class FedAVGCLient(FedTrainerClient):

diff --git a/python/fate/ml/nn/trainer/test/test_trainer.py b/python/fate/ml/nn/trainer/test/test_trainer.py
@@ -45,14 +45,14 @@ def create_ctx(local):
 
     if sys.argv[1] == "guest":
         ctx = create_ctx(guest)
-        fed_args = FedArguments(aggregate_strategy='epochs', aggregate_freq=1, aggregator='secure_aggrefgate')
+        fed_args = FedArguments(aggregate_strategy='epochs', aggregate_freq=1, aggregator='secure_aggregate')
         args = TrainingArguments(num_train_epochs=5, per_device_train_batch_size=16, logging_strategy='steps', logging_steps=5)
         trainer = FedAVGCLient(ctx=ctx, model=model, fed_args=fed_args, training_args=args, loss_fn=t.nn.BCELoss(), optimizer=t.optim.SGD(model.parameters(), lr=0.01), train_set=ds)
         trainer.train()
 
     elif sys.argv[1] == "host":
         ctx = create_ctx(host)
-        fed_args = FedArguments(aggregate_strategy='epochs', aggregate_freq=1, aggregator='plaintext')
+        fed_args = FedArguments(aggregate_strategy='epochs', aggregate_freq=1, aggregator='secure_aggregate')
         args = TrainingArguments(num_train_epochs=5, per_device_train_batch_size=16)
         trainer = FedAVGCLient(ctx=ctx, model=model, fed_args=fed_args, training_args=args, loss_fn=t.nn.BCELoss(), optimizer=t.optim.SGD(model.parameters(), lr=0.01), train_set=ds)
         trainer.train()

diff --git a/python/fate/ml/utils/_optimizer.py b/python/fate/ml/utils/_optimizer.py
@@ -346,7 +346,6 @@ def optimizer_factory(model_parameter, optimizer_type, optim_params):
 
 def lr_scheduler_factory(optimizer, method, scheduler_param):
     scheduler_method = method
-
     if scheduler_method == 'constant':
         return torch.optim.lr_scheduler.ConstantLR(
             optimizer, **scheduler_param)