update mpc nn

Signed-off-by: sagewe <wbwmat@gmail.com>
FederatedAI · Nov 21, 2023 · 16df77c · 16df77c
1 parent b8bfade
commit 16df77c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/python/fate/arch/protocol/mpc/nn/sshe/sa_layer.py b/python/fate/arch/protocol/mpc/nn/sshe/sa_layer.py
@@ -41,7 +41,6 @@ def forward(ctx, input, aggregator: "SSHEAggregator"):
     @staticmethod
     def backward(ctx: Any, *grad_outputs: Any) -> Any:
         aggregator: "SSHEAggregator" = ctx.aggregator
-        aggregator.ctx.mpc.info(f"grad_outputs={grad_outputs}", dst=[0, 1])
         ha = ctx.saved_tensors[0] if aggregator.ctx.rank == aggregator.rank_a else None
         hb = ctx.saved_tensors[0] if aggregator.ctx.rank == aggregator.rank_b else None
         dz = grad_outputs[0] if aggregator.ctx.rank == aggregator.rank_b else None

diff --git a/python/fate/ml/mpc/sshe_nn.py b/python/fate/ml/mpc/sshe_nn.py
@@ -29,7 +29,7 @@ def validate_run(self, ctx: Context):
         hb = torch.rand(num_samples, in_features_b, requires_grad=True, generator=torch.Generator().manual_seed(1))
         h = ctx.mpc.cond_call(lambda: ha, lambda: hb, dst=0)
 
-        generator = torch.Generator().manual_seed(0)
+        # generator = torch.Generator().manual_seed(0)
         layer = SSHEAggregatorLayer(
             ctx,
             in_features_a=in_features_a,
@@ -38,7 +38,7 @@ def validate_run(self, ctx: Context):
             rank_a=0,
             rank_b=1,
             lr=lr,
-            generator=generator,
+            # generator=generator,
         )
         z = layer(h)
         ctx.mpc.info(f"forward={z}", dst=[1])