MEDfl 2.0.5.dev2__tar.gz → 2.0.5.dev3__tar.gz

{medfl-2.0.5.dev2 → medfl-2.0.5.dev3}/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 → medfl-2.0.5.dev3}/MEDfl.egg-info/PKG-INFO

@@ -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 → medfl-2.0.5.dev3}/PKG-INFO

@@ -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 → medfl-2.0.5.dev3}/notebooks/data/eicu_sapsii_data.csv

@@ -1,4 +1,4 @@
 [diffend] Oversized file quarantined before diffing.
-name: medfl-2.0.5.dev2/notebooks/data/eicu_sapsii_data.csv
+name: medfl-2.0.5.dev3/notebooks/data/eicu_sapsii_data.csv
 size: 48904729 bytes
 sha256: 79203f34470b04f3b2e8cc5b3374b48e54f8a6bb6f484f259889adf7c7dff77d

{medfl-2.0.5.dev2 → medfl-2.0.5.dev3}/setup.py

@@ -14,7 +14,7 @@ with open('requirements.txt') as f:
 
 setup(
     name="MEDfl",
-    version="2.0.5.dev2",
+    version="2.0.5.dev3",
     author="MEDomics consortium",
     author_email="medomics.info@gmail.com",
     description="Python Open-source package for simulating federated learning and differential privacy",