From f48e0f942032ac516b93e9a511c4e1e32e89deae Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:06:08 +0100 Subject: [PATCH 1/8] Fix MC-dropout metrics for binary NN models --- flcore/models/nn/basic_nn.py | 7 ++++++- flcore/models/nn/client.py | 15 ++++++++------- inference.py | 13 ++++++------- 3 files changed, 20 insertions(+), 15 deletions(-) diff --git a/flcore/models/nn/basic_nn.py b/flcore/models/nn/basic_nn.py index bead24e..e71ba7c 100644 --- a/flcore/models/nn/basic_nn.py +++ b/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) diff --git a/flcore/models/nn/client.py b/flcore/models/nn/client.py index a011248..df33737 100644 --- a/flcore/models/nn/client.py +++ b/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(): diff --git a/inference.py b/inference.py index 995b525..b22e500 100644 --- a/inference.py +++ b/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): From 4f741d4c6e6ac4c5f8854fe136dd2adc812e1a02 Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:06:16 +0100 Subject: [PATCH 2/8] Save dropout_p and mc_samples in NN model metadata --- flcore/models/nn/client.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flcore/models/nn/client.py b/flcore/models/nn/client.py index df33737..c82d575 100644 --- a/flcore/models/nn/client.py +++ b/flcore/models/nn/client.py @@ -320,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"], From 8ff11cae5be3b1015831caae6fb0d1a35b140299 Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:06:16 +0100 Subject: [PATCH 3/8] Keep metadata n_out for classification targets --- flcore/datasets.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/flcore/datasets.py b/flcore/datasets.py index 7ca1372..955ffdc 100644 --- a/flcore/datasets.py +++ b/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] From c043ceb62a86c27b61fee6d6ff426885187d5938 Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:06:16 +0100 Subject: [PATCH 4/8] Aggregate eval metrics in UncertaintyWeighted strategy --- flcore/models/nn/server.py | 3 +++ server_cmd.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/flcore/models/nn/server.py b/flcore/models/nn/server.py index f8f8bee..a1ee78b 100644 --- a/flcore/models/nn/server.py +++ b/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"]) diff --git a/server_cmd.py b/server_cmd.py index 9b17624..facde77 100644 --- a/server_cmd.py +++ b/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 From 4ee85f7bdfaf9031c0251270dc79d404e268e4d6 Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:06:17 +0100 Subject: [PATCH 5/8] Use FedAvg for UncertaintyWeighted on regression --- flcore/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/flcore/utils.py b/flcore/utils.py index a719da5..31d5949 100644 --- a/flcore/utils.py +++ b/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") From 9fb4267491974dddde160153fe0b23b78258d2c5 Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:06:17 +0100 Subject: [PATCH 6/8] Respect accept_failures in UncertaintyWeightedFedAvg --- flcore/models/nn/FedCustomAggregator.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/flcore/models/nn/FedCustomAggregator.py b/flcore/models/nn/FedCustomAggregator.py index e5d9ee5..e5f51c7 100644 --- a/flcore/models/nn/FedCustomAggregator.py +++ b/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: From ecbb5471b80716c64e4feb68b682f256361f1f41 Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:06:17 +0100 Subject: [PATCH 7/8] Test NN with UncertaintyWeighted strategy --- tests/fl_runner.py | 5 +++-- tests/test_federated_training.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/fl_runner.py b/tests/fl_runner.py index 56a6a60..088c2a2 100644 --- a/tests/fl_runner.py +++ b/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", diff --git a/tests/test_federated_training.py b/tests/test_federated_training.py index dcd02b2..5c792e6 100644 --- a/tests/test_federated_training.py +++ b/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 From 3ccbb14959579b264890f06190f6ebe1609e129b Mon Sep 17 00:00:00 2001 From: faildeny Date: Thu, 29 Oct 2026 13:13:11 +0100 Subject: [PATCH 8/8] Run tests --- .github/workflows/tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 8d31f9d..0985075 100644 --- a/.github/workflows/tests.yml +++ b/.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: