fix guest gradient computation(#4659)

Signed-off-by: Yu Wu <yolandawu131@gmail.com>
FederatedAI · Jul 12, 2023 · 37a0809 · 37a0809
1 parent e803cdc
commit 37a0809
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/python/fate/ml/glm/hetero/coordinated_linr/guest.py b/python/fate/ml/glm/hetero/coordinated_linr/guest.py
@@ -155,7 +155,7 @@ def fit_model(self, ctx, train_data, validate_data=None):
                     batch_ctx.arbiter.put(loss=loss)
 
                 # gradient
-                g = 1 / h * X.T @ d
+                g = 1 / h * torch.matmul(X.T, d)
                 g = self.optimizer.add_regular_to_grad(g, w, self.init_param.get("fit_intercept"))
                 batch_ctx.arbiter.put("g_enc", g)
                 g = batch_ctx.arbiter.get("g")

diff --git a/python/fate/ml/glm/hetero/coordinated_lr/guest.py b/python/fate/ml/glm/hetero/coordinated_lr/guest.py
@@ -218,6 +218,7 @@ def fit_single_model(self, ctx: Context, train_data, validate_data=None):
                     loss -= 0.5 / h * torch.matmul(Y.T, Xw_h)
                     loss += 0.25 / h * torch.matmul(Xw.T, Xw_h)
                 if weight:
+                    logger.info(f"weight: {weight.tolist()}")
                     d = d * weight
                 batch_ctx.hosts.put(d=d)
 
@@ -232,7 +233,7 @@ def fit_single_model(self, ctx: Context, train_data, validate_data=None):
                     batch_ctx.arbiter.put(loss=loss)
 
                 # gradient
-                g = 1 / h * X.T @ d
+                g = 1 / h * torch.matmul(X.T, d)
                 g = self.optimizer.add_regular_to_grad(g, w, self.init_param.get("fit_intercept"))
                 batch_ctx.arbiter.put("g_enc", g)
                 g = batch_ctx.arbiter.get("g")