MEDfl 2.0.5.dev2__py3-none-any.whl → 2.0.5.dev3__py3-none-any.whl

MEDfl/rw/model.py

@@ -1,11 +1,8 @@
 import math
-from typing import List, Optional, Literal, Callable, Tuple
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-Task = Literal["binary", "multiclass", "multilabel", "regression"]
-
 _ACTS = {
     "relu": nn.ReLU,
     "gelu": nn.GELU,
@@ -16,7 +13,7 @@ _ACTS = {
     "identity": nn.Identity,
 }
 
-def _make_activation(name: str) -> nn.Module:
+def _make_activation(name):
     if name not in _ACTS:
         raise ValueError(f"Unsupported activation '{name}'. Choose from {list(_ACTS)}")
     return _ACTS[name]()
@@ -35,16 +32,16 @@ class Net(nn.Module):
 
     def __init__(
         self,
-        input_dim: int,
-        hidden_dims: List[int] = [64, 32],
-        activation: str = "relu",
-        batchnorm: bool = True,
-        dropout: Optional[float | List[float]] = 0.0,
-        task: Task = "binary",
-        num_classes: Optional[int] = None,
-        output_bias: bool = True,
-        return_logits: bool = True,
-        weight_init: Literal["kaiming", "xavier", "none"] = "kaiming",
+        input_dim,
+        hidden_dims=[64, 32],
+        activation="relu",
+        batchnorm=True,
+        dropout=0.0,
+        task="binary",
+        num_classes=None,
+        output_bias=True,
+        return_logits=True,
+        weight_init="kaiming",
     ):
         super().__init__()
         self.task = task
@@ -78,10 +75,9 @@ class Net(nn.Module):
         else:
             raise ValueError("dropout must be float|list[float]|None")
 
-        layers = []
-        in_dim = input_dim
         self.blocks = nn.ModuleList()
-        for i, (h, p) in enumerate(zip(hidden_dims, dropout)):
+        in_dim = input_dim
+        for h, p in zip(hidden_dims, dropout):
             block = nn.ModuleDict({
                 "lin": nn.Linear(in_dim, h, bias=True),
                 "bn": nn.BatchNorm1d(h) if batchnorm else nn.Identity(),
@@ -98,11 +94,15 @@ class Net(nn.Module):
             self.apply(lambda m: self._init_weights(m, scheme=weight_init, activation=activation))
 
     @staticmethod
-    def _init_weights(m: nn.Module, scheme: str, activation: str):
+    def _init_weights(m, scheme, activation):
         if isinstance(m, nn.Linear):
             if scheme == "kaiming":
                 nonlinearity = "leaky_relu" if activation == "leaky_relu" else "relu"
-                nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5) if activation == "leaky_relu" else 0, nonlinearity=nonlinearity)
+                nn.init.kaiming_uniform_(
+                    m.weight,
+                    a=math.sqrt(5) if activation == "leaky_relu" else 0,
+                    nonlinearity=nonlinearity,
+                )
             elif scheme == "xavier":
                 nn.init.xavier_uniform_(m.weight)
             if m.bias is not None:
@@ -110,7 +110,7 @@ class Net(nn.Module):
                 bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
                 nn.init.uniform_(m.bias, -bound, bound)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x):
         for blk in self.blocks:
             x = blk["lin"](x)
             x = blk["bn"](x)
@@ -120,7 +120,7 @@ class Net(nn.Module):
         return logits if self.return_logits else self._apply_output_activation(logits)
 
     # Inference helpers
-    def _apply_output_activation(self, logits: torch.Tensor) -> torch.Tensor:
+    def _apply_output_activation(self, logits):
         if self.task == "binary":
             return torch.sigmoid(logits)               # (B, 1)
         elif self.task == "multiclass":
@@ -128,14 +128,15 @@ class Net(nn.Module):
         elif self.task == "multilabel":
             return torch.sigmoid(logits)               # (B, L)
         elif self.task == "regression":
-            return logits                               # raw regression output
+            return logits                              # raw regression output
         else:
             raise RuntimeError("Invalid task")
 
     @torch.no_grad()
-    def predict(self, x: torch.Tensor) -> torch.Tensor:
-        out = self.forward(x)                                  # always run forward
+    def predict(self, x):
+        out = self.forward(x)
         probs = self._apply_output_activation(out) if self.return_logits else out
+
         if self.task == "binary":
             return (probs >= 0.5).long()
         elif self.task == "multiclass":
@@ -144,4 +145,3 @@ class Net(nn.Module):
             return (probs >= 0.5).long()
         elif self.task == "regression":
             return probs
-

{medfl-2.0.5.dev2.dist-info → medfl-2.0.5.dev3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: MEDfl
-Version: 2.0.5.dev2
+Version: 2.0.5.dev3
 Summary: Python Open-source package for simulating federated learning and differential privacy
 Home-page: https://github.com/MEDomics-UdeS/MEDfl
 Author: MEDomics consortium

{medfl-2.0.5.dev2.dist-info → medfl-2.0.5.dev3.dist-info}/RECORD

@@ -21,7 +21,7 @@ MEDfl/NetManager/network.py,sha256=5t705fzWc-BRg-QPAbAcDv5ckDGzsPwj_Q5V0iTgkx0,6
 MEDfl/NetManager/node.py,sha256=t90QuYZ8M1X_AG1bwTta0CnlOuodqkmpVda2K7NOgHc,6542
 MEDfl/rw/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 MEDfl/rw/client.py,sha256=Ufa-VCbP83adEs7tV1ltDYakBV2Ge5GCPhjIJ7HsuJ8,17804
-MEDfl/rw/model.py,sha256=-YRBO5ZX_Ahp95lTisCUejlWJHS1Ly5TnRMp4LhA0Do,5759
+MEDfl/rw/model.py,sha256=IvC5a1ADyn7mKeSBq1Lk_UnkSJOCErPCmpEPSOFxHpk,5343
 MEDfl/rw/server.py,sha256=TmcWUwMJ5BG7owqIVTtTC6w8bR02SESRT9lXh7BqlOg,4986
 MEDfl/rw/strategy.py,sha256=rrRUvuFxL4rdtjN-q-GgB_QQ4ntH9kSyIbTQtiASjVI,13779
 MEDfl/scripts/__init__.py,sha256=Pq1weevsPaU7MRMHfBYeyT0EOFeWLeVM6Y1DVz6jw1A,48
@@ -29,8 +29,8 @@ MEDfl/scripts/base.py,sha256=QrmG7gkiPYkAy-5tXxJgJmOSLGAKeIVH6i0jq7G9xnA,752
 MEDfl/scripts/create_db.py,sha256=MnFtZkTueRZ-3qXPNX4JsXjOKj-4mlkxoRhSFdRcvJw,3817
 alembic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 alembic/env.py,sha256=-aSZ6SlJeK1ZeqHgM-54hOi9LhJRFP0SZGjut-JnY-4,1588
-medfl-2.0.5.dev2.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-medfl-2.0.5.dev2.dist-info/METADATA,sha256=dZQUWWXwA-R_MHYDgp24UvIWeUmUf2RxNe4BTorptME,6914
-medfl-2.0.5.dev2.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-medfl-2.0.5.dev2.dist-info/top_level.txt,sha256=dIL9X8HOFuaVSzpg40DVveDPrymWRoShHtspH7kkjdI,14
-medfl-2.0.5.dev2.dist-info/RECORD,,
+medfl-2.0.5.dev3.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+medfl-2.0.5.dev3.dist-info/METADATA,sha256=jqEhW5WJalcs9sBxuN27oY5_e1yjwIRUuvY7pzrSpII,6914
+medfl-2.0.5.dev3.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+medfl-2.0.5.dev3.dist-info/top_level.txt,sha256=dIL9X8HOFuaVSzpg40DVveDPrymWRoShHtspH7kkjdI,14
+medfl-2.0.5.dev3.dist-info/RECORD,,