Bugfix: Multiclass in random forest and proper model saving

.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/models/random_forest/client.py

@@ -113,9 +113,6 @@ def fit(self, ins: FitIns):  # , parameters, config type: ignore
             params = utils.get_model_parameters(self.model)
             parameters_updated = serialize_RF(params)
 
-            if self.round % self.config["save_every_n_rounds"] == 0:
-                self.save_model()
-
             self.round += 1
 
             # Build and return response
@@ -140,6 +137,11 @@ def evaluate(self, ins: EvaluateIns):  # , parameters, config type: ignore
             #Deserialize to get the real parameters
             parameters = deserialize_RF(parameters)
             utils.set_model_params(self.model, parameters)
+
+            # Save the aggregated model received from the server
+            rnd = self.round - 1
+            if rnd % self.config["save_every_n_rounds"] == 0 or rnd == self.config["num_rounds"] - 1:
+                self.save_model(rnd)
 
             ## AQUI TAMBIEN TENDRIAMOS QUE ADAPTAR PARA REGRESOR/CLASIFICADOR
             if self.config["task"] == "classification":
@@ -210,12 +212,12 @@ def evaluate(self, ins: EvaluateIns):  # , parameters, config type: ignore
             log_detailed_error("Model Evaluation (Local Validation)", e, config=getattr(self, "config", None), X=getattr(self, "X_test", None), y=getattr(self, "y_test", None))
             raise e
 
-    def save_model(self):
+    def save_model(self, rnd):
         save_path = Path(self.config["experiment_dir"]) / "models"
         save_path.mkdir(parents=True, exist_ok=True)
-        is_final = self.round == self.config["num_rounds"] - 1
+        is_final = rnd == self.config["num_rounds"] - 1
 
-        model_name = f"{self.config['model']}_{self.config['task']}_round_{self.round}"
+        model_name = f"{self.config['model']}_{self.config['task']}_round_{rnd}"
 
         with open(self.config["metadata_file"], "r") as f:
             data_metadata = json.load(f)
@@ -251,7 +253,6 @@ def save_model(self):
             elif label in all_features_meta:
                 outcomes_meta[label] = all_features_meta[label]
 
-        self.model.n_outputs_ = self.config["n_out"]
         self.model.n_features_in_ = self.config["n_feats"]
         n_classes = self.config.get("n_classes")
         if n_classes is not None:
@@ -269,7 +270,7 @@ def save_model(self):
             "features_meta": features_meta,
             "outcomes_meta": outcomes_meta,
             "is_final": is_final,
-            "round": self.round,
+            "round": rnd,
             "model_name": model_name,
             "sklearn_version": sklearn.__version__,
             "imblearn_version": getattr(imblearn, "__version__", None) if imblearn else None,
@@ -290,4 +291,4 @@ def save_model(self):
 def get_client(config,data) -> fl.client.Client:
     return MnistClient(data, config)
     # # Start Flower client
-    # fl.client.start_numpy_client(server_address="0.0.0.0:8080", client=MnistClient())
+    # fl.client.start_numpy_client(server_address="0.0.0.0:8080", client=MnistClient())

tests/test_federated_training.py

@@ -86,11 +86,7 @@ def case(model, scenario, client_extra=(), id_suffix="", bug=None):
 # Tree ensembles and neural network
 CASES += [
     case("random_forest", "binary"),
-    case(
-        "random_forest",
-        "multiclass",
-        bug="random_forest client.evaluate crashes on multiclass probabilities",
-    ),
+    case("random_forest", "multiclass"),
     case("random_forest", "regression"),
 ]
 wrf_bug = "weighted_random_forest server reads config['weighted_random_forest'], never set by the CLI"