Fixes in Nn uncertainty model

.github/workflows/tests.yml

@@ -4,7 +4,7 @@ on:
   push:
     branches: [main]
   pull_request:
-    branches: [main, cmd_line]
+    branches: [main, cmd_line, tests_and_CI]
   workflow_dispatch:
 
 concurrency:

flcore/datasets.py

@@ -861,10 +861,10 @@ def load_dt4h(config):
     X = dat[train_labels]
     y = dat[target_labels].iloc[:, 0]
 
-    # Calculate n_out dynamically
-    if config.get("task") == "multiclass":
-        config["n_out"] = len(np.unique(y))
-    elif config.get("task") == "classification" and len(np.unique(y)) > 2:
+    # n_out comes from the shared metadata (CheckClientConfig) so every site
+    # agrees on it; the local data may lack some classes. Only fall back to
+    # the local count when the metadata could not provide it.
+    if not config.get("n_out") and config.get("task") in ("classification", "multiclass"):
         config["n_out"] = len(np.unique(y))
 
     X_train = X[:split_idx]

flcore/models/nn/FedCustomAggregator.py

@@ -34,6 +34,9 @@ def __init__(self, epsilon: float = 0.1, gamma: float = 1.0, **kwargs):
     def aggregate_fit(self, server_round, results, failures):
         if not results:
             return None, {}
+        # As in FedAvg: do not aggregate if there were failures and they are not accepted
+        if not self.accept_failures and failures:
+            return None, {}
 
         ndarrays_list, raw_weights = [], []
         for _, fit_res in results:

flcore/models/nn/basic_nn.py

@@ -33,7 +33,12 @@ def predict_proba_mc(self, x, T: int = 20):
         probs = []
         for _ in range(T):
             logits = self(x)
-            probs.append(F.softmax(logits, dim=-1))
+            if logits.shape[-1] == 1:
+                # Binary (BCEWithLogitsLoss): softmax of a single logit is always 1
+                p1 = torch.sigmoid(logits)
+                probs.append(torch.cat([1 - p1, p1], dim=-1))
+            else:
+                probs.append(F.softmax(logits, dim=-1))
 
         probs = torch.stack(probs, dim=0)  # [T, B, C]
         mean = probs.mean(dim=0)

flcore/models/nn/client.py

@@ -238,14 +238,15 @@ def evaluate(self, parameters, config):
             self.model.eval()
             y_true = self.y_test.cpu().numpy()
 
-            if self.config["dropout_p"] > 0.0:
-                metrics = uncertainty_metrics(
+            y_pred = self._predict(self.X_test)
+            metrics = calculate_metrics(y_true, y_pred, self.config)
+
+            # MC Dropout only makes sense with class probabilities
+            if self.config["dropout_p"] > 0.0 and self.config["task"] == "classification":
+                metrics.update(uncertainty_metrics(
                     self.model, self.val_loader, device=self.device, T=int(self.config["T"])
-                )
-                self.model.eval()  # por si uncertainty_metrics deja dropout activo
-            else:
-                y_pred = self._predict(self.X_test)
-                metrics = calculate_metrics(y_true, y_pred, self.config)
+                ))
+                self.model.eval()  # uncertainty_metrics leaves dropout active
 
             total_loss, total = 0.0, 0
             with torch.no_grad():
@@ -319,6 +320,8 @@ def save_model(self):
             "task": self.config["task"],
             "n_out": self.config["n_out"],
             "n_feats": self.config["n_feats"],
+            "dropout_p": self.config["dropout_p"],
+            "mc_samples": int(self.config["T"]),
             "model_type": self.config["model"],
             "feature_names": self.config["train_labels"],
             "target_names": self.config["target_labels"],

flcore/models/nn/server.py

@@ -88,6 +88,9 @@ def get_server_and_strategy(config):
         min_available_clients = config["min_available_clients"])
     elif config["strategy"] == "UncertaintyWeighted":
         strategy = UncertaintyWeightedFedAvg(
+        epsilon = config.get("uw_epsilon", 0.1),
+        gamma = config.get("uw_gamma", 1.0),
+        evaluate_metrics_aggregation_fn=metrics,
         min_fit_clients = config["min_fit_clients"],
         min_evaluate_clients = config["min_evaluate_clients"],
         min_available_clients = config["min_available_clients"])

flcore/utils.py

@@ -348,7 +348,11 @@ def CheckServerConfig(config):
     """
     if config["strategy"] == "UncertaintyWeighted":
         if config["model"] == "nn":
-            pass
+            if config["task"] != "classification":
+                # Clients only report entropy for classification
+                print("UncertaintyWeighted needs class probabilities (classification task)")
+                print("Changing strategy to FedAvg")
+                config["strategy"] = "FedAvg"
         else:
            print("UncertaintyWeighted is only available for NN")
            print("Changing strategy to FedAvg")

inference.py

@@ -126,13 +126,12 @@ def load_model(model_dir, model=None, task=None, prefer_final=True, mlp_cls=None
 
         class NnWrapperTorch:
             """
-            task/n_out se guardan aquí, no en el modelo: BasicNN no los expone
-            como atributos. La activación de salida se decide igual que en
-            FlowerClient.evaluate() (sigmoid para binario, softmax para
-            multiclase, nada para regresión) — no se reutiliza
-            BasicNN.predict_proba_mc porque ese método aplica softmax
-            incondicionalmente, lo cual es incorrecto para el caso binario
-            (entrenado con BCEWithLogitsLoss/sigmoid, no con softmax).
+            task/n_out are stored here, not on the model: BasicNN does not
+            expose them as attributes. The output activation matches
+            FlowerClient.evaluate() (sigmoid for binary, softmax for
+            multiclass, none for regression). BasicNN.predict_proba_mc is not
+            reused because it returns class probabilities, which do not apply
+            to regression.
             """
 
             def __init__(self, m, task, n_out, n_samples=20):

server_cmd.py

@@ -39,6 +39,8 @@
     parser.add_argument("--dropout_percentage", type=float, default=0.0, help="Ratio of dropout nodes")
     parser.add_argument("--checkpoint_selection_metric", type=str, default="precision", help="Metric used for checkpoints")
     parser.add_argument("--metrics_aggregation", type=str, default="weighted_average",  help="Metrics")
+    parser.add_argument("--uw_epsilon", type=float, default=0.1, help="UncertaintyWeighted: smooths weights of clients with entropy ~0")
+    parser.add_argument("--uw_gamma", type=float, default=1.0, help="UncertaintyWeighted: 0 => FedAvg, larger => more weight to low entropy")
     parser.add_argument("--experiment_name", type=str, default="experiment_1", help="Experiment directory")
 
     # Client connection wait settings

tests/fl_runner.py

@@ -99,13 +99,14 @@ def _attempt(
     num_rounds=2,
     extra_args=(),
     client_extra_args=(),
+    server_extra_args=(),
     timeout=300,
     seed=42,
 ):
     """Train ``model`` for ``num_rounds`` with ``num_clients`` clients.
 
     ``extra_args`` go to both server and clients, ``client_extra_args`` only to
-    the clients. Never raises on a failed training: returns a
+    the clients and ``server_extra_args`` only to the server. Never raises on a failed training: returns a
     :class:`FederationResult` so the caller decides what to assert.
     """
     sandbox = Path(sandbox)
@@ -130,7 +131,7 @@ def _attempt(
         *_opt("levelOfDetail", "DecisionTree"),
         *extra_args,
     ]
-    server_cmd = [sys.executable, "server_cmd.py", *common]
+    server_cmd = [sys.executable, "server_cmd.py", *common, *server_extra_args]
     client_cmd = [
         sys.executable,
         "client_cmd.py",

tests/test_federated_training.py

@@ -134,3 +134,33 @@ def test_federated_training(
         + result.describe()
     )
     assert all(math.isfinite(x) for x in losses), f"non-finite losses: {losses}"
+
+
+@pytest.mark.parametrize("scenario", ["binary", "multiclass", "regression"])
+def test_nn_uncertainty_weighted(scenario, synthetic_data_dir, tmp_path):
+    """NN with MC dropout and the uncertainty-weighted aggregation strategy."""
+    task, target_labels = TARGETS[scenario]
+    result = run_federation(
+        model="nn",
+        task=task,
+        data_path=synthetic_data_dir,
+        sandbox=tmp_path / "sandbox",
+        target_labels=target_labels,
+        num_clients=NUM_CLIENTS,
+        num_rounds=NUM_ROUNDS,
+        client_extra_args=["--dropout_p", "0.2", "--T", "5"],
+        server_extra_args=["--strategy", "UncertaintyWeighted"],
+    )
+    assert result.ok, result.describe()
+
+    losses = result.history.get("losses_distributed", [])
+    assert len(losses) == NUM_ROUNDS, result.describe()
+    assert all(math.isfinite(x) for x in losses), f"non-finite losses: {losses}"
+    # Regression falls back to FedAvg; either way clients must not be
+    # aggregated with a missing entropy.
+    assert "sin 'entropy'" not in result.server_log, result.describe()
+
+    metrics = result.history.get("metrics_distributed", {})
+    assert metrics, "server did not aggregate evaluation metrics"
+    if task == "classification":
+        assert "entropy" in metrics and "val_accuracy" in metrics, metrics