From 95c0c823c1524d5b67a6ce6d3775924ff9f425aa Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Thu, 29 Oct 2026 22:35:15 -0400 Subject: [PATCH 1/3] add more metrics --- scripts/run_experiments.py | 5 -- tests/scripts/test_run_experiments.py | 3 +- .../training_plans/test_tabular_classifier.py | 48 +++++++++++++++++++ .../training_plans/test_tabular_regressor.py | 25 ++++++++++ training_plans/tabular_classifier.py | 17 +++++++ training_plans/tabular_regressor.py | 14 +++++- 6 files changed, 105 insertions(+), 7 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 9bbdffa..2ab4dd6 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -16,7 +16,6 @@ import numpy as np from declearn.optimizer.modules import ScaffoldServerModule -from fedbiomed.common.metrics import MetricTypes from fedbiomed.common.optimizers.optimizer import Optimizer from fedbiomed.common.training_args import TrainingArgs from fedbiomed.common.training_plans import BaseTrainingPlan @@ -48,7 +47,6 @@ class TaskConfig: training_plan_module: str training_plan_class: str target_column: str - test_metric: MetricTypes # whether the model intercept is initialized to the global target mean with_target_mean: bool = False extra_model_args: dict = field(default_factory=dict) @@ -97,7 +95,6 @@ def get_fpath_out( training_plan_module="tabular_regressor", training_plan_class="TabularRegressorTrainingPlan", target_column="nb:Age", - test_metric=MetricTypes.MEAN_SQUARE_ERROR, with_target_mean=True, secagg_clipping_range=100, ), @@ -105,7 +102,6 @@ def get_fpath_out( training_plan_module="tabular_classifier", training_plan_class="TabularClassifierTrainingPlan", target_column="nb:Sex", - test_metric=MetricTypes.ACCURACY, extra_model_args={"n_classes": 2}, ), } @@ -469,7 +465,6 @@ def _run_task( "loader_args": {"batch_size": batch_size}, "num_updates": n_updates, "test_ratio": test_ratio, - "test_metric": config.test_metric, "test_on_global_updates": True, "test_on_local_updates": True, "optimizer_args": _build_optimizer_args( diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 2c21d89..119b16e 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -577,7 +577,8 @@ def test_run_task( ] exp.set_model_args.assert_called_once_with(model_args) training_args = json_data["settings"]["training_args"] - assert training_args["test_metric"] == "MEAN_SQUARE_ERROR" + # the metrics are computed by the training plan testing_step + assert "test_metric" not in training_args # JSON object keys are strings assert json_data["results"] == json.loads( json.dumps( diff --git a/tests/training_plans/test_tabular_classifier.py b/tests/training_plans/test_tabular_classifier.py index 9601be6..c4cef66 100644 --- a/tests/training_plans/test_tabular_classifier.py +++ b/tests/training_plans/test_tabular_classifier.py @@ -5,6 +5,7 @@ import numpy as np import pytest from fedbiomed.common.dataset import Dataset +from fedbiomed.common.training_args import TrainingArgs from fedbiomed.common.training_plans import BaseTrainingPlan @@ -106,3 +107,50 @@ def test_no_training_data(self, plan_cls: type[BaseTrainingPlan]): plan.set_data_loaders(None, test_data_loader) assert plan.training_data_loader is None assert plan.testing_data_loader is test_data_loader + + @pytest.mark.parametrize("use_source", [False, True]) + def test_testing_step( + self, + plan_cls: type[BaseTrainingPlan], + model_args: dict, + from_source: Callable[[BaseTrainingPlan], type[BaseTrainingPlan]], + use_source: bool, + ): + plan = _fitted_plan(plan_cls, model_args, from_source if use_source else None) + # decision scores: [-2, -1, 1, 2], predictions: [0, 0, 1, 1] + data = np.array([[-2.0, 5.0], [-1.0, 5.0], [1.0, 5.0], [2.0, 5.0]]) + target = np.array([[0], [1], [1], [1]]) + metrics = plan.testing_step(data, target) + assert metrics.keys() == {"ACCURACY", "BALANCED_ACCURACY", "AUROC"} + np.testing.assert_allclose(metrics["ACCURACY"], 0.75) + # (1 + 2 / 3) / 2 + np.testing.assert_allclose(metrics["BALANCED_ACCURACY"], 5 / 6) + np.testing.assert_allclose(metrics["AUROC"], 1.0) + + @pytest.mark.filterwarnings("ignore:A single label was found") + def test_testing_step_single_class( + self, plan_cls: type[BaseTrainingPlan], model_args: dict + ): + plan = _fitted_plan(plan_cls, model_args) + data = np.array([[1.0, 5.0], [2.0, 5.0]]) + metrics = plan.testing_step(data, np.array([[1], [1]])) + np.testing.assert_allclose(metrics["ACCURACY"], 1.0) + assert np.isnan(metrics["AUROC"]) + + +def _fitted_plan( + plan_cls: type[BaseTrainingPlan], + model_args: dict, + from_source: Callable[[BaseTrainingPlan], type[BaseTrainingPlan]] | None = None, +) -> BaseTrainingPlan: + """Set up a training plan with the decision function x[0].""" + plan = plan_cls() + if from_source is not None: + plan = from_source(plan)() + plan.post_init(model_args, TrainingArgs({}, only_required=False)) + plan.set_model_params( + {"coef_": np.array([[1.0, 0.0]]), "intercept_": np.array([0.0])} + ) + # set by Fed-BioMed before testing + plan.model().classes_ = np.array([0, 1]) + return plan diff --git a/tests/training_plans/test_tabular_regressor.py b/tests/training_plans/test_tabular_regressor.py index 01c30dc..c397e07 100644 --- a/tests/training_plans/test_tabular_regressor.py +++ b/tests/training_plans/test_tabular_regressor.py @@ -124,3 +124,28 @@ def test_no_training_data(self, plan_cls: type[BaseTrainingPlan]): plan.set_data_loaders(None, test_data_loader) assert plan.training_data_loader is None assert plan.testing_data_loader is test_data_loader + + @pytest.mark.parametrize("use_source", [False, True]) + def test_testing_step( + self, + plan_cls: type[BaseTrainingPlan], + model_args: dict, + from_source: Callable[[BaseTrainingPlan], type[BaseTrainingPlan]], + use_source: bool, + ): + plan = plan_cls() + if use_source: + plan = from_source(plan)() + plan.post_init(model_args, TrainingArgs({}, only_required=False)) + plan.set_model_params( + {"coef_": np.array([1.0, 0.0]), "intercept_": np.array([0.0])} + ) + # predictions: [1, 2, 3, 4] + data = np.array([[1.0, 5.0], [2.0, 5.0], [3.0, 5.0], [4.0, 5.0]]) + target = np.array([[1.0], [2.0], [4.0], [5.0]]) + metrics = plan.testing_step(data, target) + assert metrics.keys() == {"MEAN_SQUARE_ERROR", "MEAN_ABSOLUTE_ERROR", "R2"} + np.testing.assert_allclose(metrics["MEAN_SQUARE_ERROR"], 0.5) + np.testing.assert_allclose(metrics["MEAN_ABSOLUTE_ERROR"], 0.5) + # 1 - 2 / 10 + np.testing.assert_allclose(metrics["R2"], 0.8) diff --git a/training_plans/tabular_classifier.py b/training_plans/tabular_classifier.py index fe39afa..dba1135 100644 --- a/training_plans/tabular_classifier.py +++ b/training_plans/tabular_classifier.py @@ -7,6 +7,7 @@ from fedbiomed.common.dataset import TabularDataset from fedbiomed.common.optimizers.optimizer import Optimizer from fedbiomed.common.training_plans import FedSGDClassifier +from sklearn.metrics import accuracy_score, balanced_accuracy_score, roc_auc_score class TabularClassifierTrainingPlan(FedSGDClassifier): @@ -20,6 +21,8 @@ def init_dependencies(self): "from fedbiomed.common.optimizers.optimizer import Optimizer", "from declearn.optimizer.modules import ScaffoldClientModule", "from declearn.optimizer.regularizers import FedProxRegularizer", + "from sklearn.metrics import accuracy_score, balanced_accuracy_score", + "from sklearn.metrics import roc_auc_score", ] def init_optimizer(self, optimizer_args): @@ -63,3 +66,17 @@ def training_data(self): transform=lambda x: (x - mean) / std, ) return DataManager(dataset=dataset, shuffle=True) + + def testing_step(self, data, target): + target = np.ravel(target) + prediction = self.model().predict(data) + if len(np.unique(target)) < 2: + # AUROC is undefined if the split has a single class + auroc = float("nan") + else: + auroc = roc_auc_score(target, self.model().decision_function(data)) + return { + "ACCURACY": accuracy_score(target, prediction), + "BALANCED_ACCURACY": balanced_accuracy_score(target, prediction), + "AUROC": auroc, + } diff --git a/training_plans/tabular_regressor.py b/training_plans/tabular_regressor.py index 751300a..788b09f 100644 --- a/training_plans/tabular_regressor.py +++ b/training_plans/tabular_regressor.py @@ -7,6 +7,7 @@ from fedbiomed.common.dataset import TabularDataset from fedbiomed.common.optimizers.optimizer import Optimizer from fedbiomed.common.training_plans import FedSGDRegressor +from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score class TabularRegressorTrainingPlan(FedSGDRegressor): @@ -20,6 +21,7 @@ def init_dependencies(self): "from fedbiomed.common.optimizers.optimizer import Optimizer", "from declearn.optimizer.modules import ScaffoldClientModule", "from declearn.optimizer.regularizers import FedProxRegularizer", + "from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score", ] def init_optimizer(self, optimizer_args): @@ -46,7 +48,8 @@ def post_init(self, model_args, training_args, *args, **kwargs): def set_data_loaders(self, train_data_loader, test_data_loader): if train_data_loader is None: # no training data when the node only tests the model - # used in held-out node validation self.training_data_loader = None + # used in held-out node validation + self.training_data_loader = None self.testing_data_loader = test_data_loader else: super().set_data_loaders(train_data_loader, test_data_loader) @@ -68,3 +71,12 @@ def training_data(self): transform=lambda x: (x - mean) / std, ) return DataManager(dataset=dataset, shuffle=True) + + def testing_step(self, data, target): + target = np.ravel(target) + prediction = self.model().predict(data) + return { + "MEAN_SQUARE_ERROR": mean_squared_error(target, prediction), + "MEAN_ABSOLUTE_ERROR": mean_absolute_error(target, prediction), + "R2": r2_score(target, prediction), + } From 2e1d3df36a254005c0f49da9e20b85004ad3dfe5 Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Thu, 29 Oct 2026 23:11:21 -0400 Subject: [PATCH 2/3] allow setting of metrics via model args, and drop metrics that have NA values --- scripts/run_experiments.py | 8 +++- .../training_plans/test_tabular_classifier.py | 37 +++++++++++++++++-- .../training_plans/test_tabular_regressor.py | 37 ++++++++++++++++++- training_plans/tabular_classifier.py | 30 ++++++++++----- training_plans/tabular_regressor.py | 19 ++++++++-- 5 files changed, 113 insertions(+), 18 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 2ab4dd6..842fcb4 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -96,13 +96,19 @@ def get_fpath_out( training_plan_class="TabularRegressorTrainingPlan", target_column="nb:Age", with_target_mean=True, + extra_model_args={ + "metrics": ["MEAN_SQUARE_ERROR", "MEAN_ABSOLUTE_ERROR", "R2"] + }, secagg_clipping_range=100, ), "sex": TaskConfig( training_plan_module="tabular_classifier", training_plan_class="TabularClassifierTrainingPlan", target_column="nb:Sex", - extra_model_args={"n_classes": 2}, + extra_model_args={ + "n_classes": 2, + "metrics": ["ACCURACY", "BALANCED_ACCURACY", "AUROC"], + }, ), } diff --git a/tests/training_plans/test_tabular_classifier.py b/tests/training_plans/test_tabular_classifier.py index c4cef66..ebeba5c 100644 --- a/tests/training_plans/test_tabular_classifier.py +++ b/tests/training_plans/test_tabular_classifier.py @@ -116,7 +116,11 @@ def test_testing_step( from_source: Callable[[BaseTrainingPlan], type[BaseTrainingPlan]], use_source: bool, ): - plan = _fitted_plan(plan_cls, model_args, from_source if use_source else None) + plan = _fitted_plan( + plan_cls, + model_args | {"metrics": list(plan_cls.METRICS)}, + from_source if use_source else None, + ) # decision scores: [-2, -1, 1, 2], predictions: [0, 0, 1, 1] data = np.array([[-2.0, 5.0], [-1.0, 5.0], [1.0, 5.0], [2.0, 5.0]]) target = np.array([[0], [1], [1], [1]]) @@ -128,14 +132,39 @@ def test_testing_step( np.testing.assert_allclose(metrics["AUROC"], 1.0) @pytest.mark.filterwarnings("ignore:A single label was found") + @pytest.mark.filterwarnings("ignore:Only one class is present") def test_testing_step_single_class( self, plan_cls: type[BaseTrainingPlan], model_args: dict ): - plan = _fitted_plan(plan_cls, model_args) + plan = _fitted_plan(plan_cls, model_args | {"metrics": list(plan_cls.METRICS)}) data = np.array([[1.0, 5.0], [2.0, 5.0]]) metrics = plan.testing_step(data, np.array([[1], [1]])) - np.testing.assert_allclose(metrics["ACCURACY"], 1.0) - assert np.isnan(metrics["AUROC"]) + assert metrics == {"ACCURACY": 1.0, "BALANCED_ACCURACY": 1.0} + + def test_testing_step_default( + self, plan_cls: type[BaseTrainingPlan], model_args: dict + ): + plan = _fitted_plan(plan_cls, model_args) + # predictions: [0, 0, 1] + data = np.array([[-2.0, 5.0], [-1.0, 5.0], [1.0, 5.0]]) + target = np.array([[0], [1], [1]]) + metrics = plan.testing_step(data, target) + assert metrics.keys() == {"ACCURACY"} + np.testing.assert_allclose(metrics["ACCURACY"], 2 / 3) + + def test_testing_step_metrics_arg( + self, plan_cls: type[BaseTrainingPlan], model_args: dict + ): + plan = _fitted_plan(plan_cls, model_args | {"metrics": ["AUROC"]}) + data = np.array([[-1.0, 5.0], [1.0, 5.0]]) + assert plan.testing_step(data, np.array([[0], [1]])) == {"AUROC": 1.0} + + def test_unknown_metric(self, plan_cls: type[BaseTrainingPlan], model_args: dict): + with pytest.raises(ValueError, match="Unknown metrics"): + plan_cls().post_init( + model_args | {"metrics": ["AUROC", "F1"]}, + TrainingArgs({}, only_required=False), + ) def _fitted_plan( diff --git a/tests/training_plans/test_tabular_regressor.py b/tests/training_plans/test_tabular_regressor.py index c397e07..00d8a25 100644 --- a/tests/training_plans/test_tabular_regressor.py +++ b/tests/training_plans/test_tabular_regressor.py @@ -136,7 +136,10 @@ def test_testing_step( plan = plan_cls() if use_source: plan = from_source(plan)() - plan.post_init(model_args, TrainingArgs({}, only_required=False)) + plan.post_init( + model_args | {"metrics": list(plan_cls.METRICS)}, + TrainingArgs({}, only_required=False), + ) plan.set_model_params( {"coef_": np.array([1.0, 0.0]), "intercept_": np.array([0.0])} ) @@ -149,3 +152,35 @@ def test_testing_step( np.testing.assert_allclose(metrics["MEAN_ABSOLUTE_ERROR"], 0.5) # 1 - 2 / 10 np.testing.assert_allclose(metrics["R2"], 0.8) + + def test_testing_step_metrics_arg( + self, plan_cls: type[BaseTrainingPlan], model_args: dict + ): + plan = plan_cls() + plan.post_init( + model_args | {"metrics": ["MEAN_ABSOLUTE_ERROR"]}, + TrainingArgs({}, only_required=False), + ) + plan.set_model_params( + {"coef_": np.array([1.0, 0.0]), "intercept_": np.array([0.0])} + ) + metrics = plan.testing_step(np.array([[1.0, 5.0]]), np.array([[3.0]])) + assert metrics == {"MEAN_ABSOLUTE_ERROR": 2.0} + + def test_unknown_metric(self, plan_cls: type[BaseTrainingPlan], model_args: dict): + with pytest.raises(ValueError, match="Unknown metrics"): + plan_cls().post_init( + model_args | {"metrics": ["R2", "RMSE"]}, + TrainingArgs({}, only_required=False), + ) + + def test_testing_step_default( + self, plan_cls: type[BaseTrainingPlan], model_args: dict + ): + plan = plan_cls() + plan.post_init(model_args, TrainingArgs({}, only_required=False)) + plan.set_model_params( + {"coef_": np.array([1.0, 0.0]), "intercept_": np.array([0.0])} + ) + metrics = plan.testing_step(np.array([[1.0, 5.0]]), np.array([[3.0]])) + assert metrics == {"MEAN_SQUARE_ERROR": 4.0} diff --git a/training_plans/tabular_classifier.py b/training_plans/tabular_classifier.py index dba1135..6d6ddfe 100644 --- a/training_plans/tabular_classifier.py +++ b/training_plans/tabular_classifier.py @@ -13,6 +13,10 @@ class TabularClassifierTrainingPlan(FedSGDClassifier): """Predict an integer-encoded class column from standardized feature columns.""" + # allows metrics + METRICS = ("ACCURACY", "BALANCED_ACCURACY", "AUROC") + DEFAULT_METRICS = ("ACCURACY",) + def init_dependencies(self): return [ "import numpy as np", @@ -21,8 +25,7 @@ def init_dependencies(self): "from fedbiomed.common.optimizers.optimizer import Optimizer", "from declearn.optimizer.modules import ScaffoldClientModule", "from declearn.optimizer.regularizers import FedProxRegularizer", - "from sklearn.metrics import accuracy_score, balanced_accuracy_score", - "from sklearn.metrics import roc_auc_score", + "from sklearn.metrics import accuracy_score, balanced_accuracy_score, roc_auc_score", ] def init_optimizer(self, optimizer_args): @@ -40,6 +43,14 @@ def init_optimizer(self, optimizer_args): lr=optimizer_args["lr"], modules=modules, regularizers=regularizers ) + def post_init(self, model_args, training_args, *args, **kwargs): + super().post_init(model_args, training_args, *args, **kwargs) + unknown = set(model_args.get("metrics", [])) - set(self.METRICS) + if unknown: + raise ValueError( + f"Unknown metrics {sorted(unknown)}. Allowed metrics: {list(self.METRICS)}" + ) + def set_data_loaders(self, train_data_loader, test_data_loader): if train_data_loader is None: # no training data when the node only tests the model @@ -70,13 +81,14 @@ def training_data(self): def testing_step(self, data, target): target = np.ravel(target) prediction = self.model().predict(data) - if len(np.unique(target)) < 2: - # AUROC is undefined if the split has a single class - auroc = float("nan") - else: - auroc = roc_auc_score(target, self.model().decision_function(data)) - return { + metrics = { "ACCURACY": accuracy_score(target, prediction), "BALANCED_ACCURACY": balanced_accuracy_score(target, prediction), - "AUROC": auroc, + "AUROC": roc_auc_score(target, self.model().decision_function(data)), + } + # NaN cannot be saved in the results JSON + return { + name: metrics[name] + for name in self.model_args().get("metrics", self.DEFAULT_METRICS) + if not np.isnan(metrics[name]) } diff --git a/training_plans/tabular_regressor.py b/training_plans/tabular_regressor.py index 788b09f..9058ef7 100644 --- a/training_plans/tabular_regressor.py +++ b/training_plans/tabular_regressor.py @@ -13,6 +13,10 @@ class TabularRegressorTrainingPlan(FedSGDRegressor): """Predict a numeric column from standardized feature columns.""" + # allowed metrics + METRICS = ("MEAN_SQUARE_ERROR", "MEAN_ABSOLUTE_ERROR", "R2") + DEFAULT_METRICS = ("MEAN_SQUARE_ERROR",) + def init_dependencies(self): return [ "import numpy as np", @@ -41,6 +45,11 @@ def init_optimizer(self, optimizer_args): def post_init(self, model_args, training_args, *args, **kwargs): super().post_init(model_args, training_args, *args, **kwargs) + unknown = set(model_args.get("metrics", [])) - set(self.METRICS) + if unknown: + raise ValueError( + f"Unknown metrics {sorted(unknown)}. Allowed metrics: {list(self.METRICS)}" + ) # start from the global mean of the target instead of 0 if "target_mean" in model_args: self.model().intercept_ = np.array([float(model_args["target_mean"])]) @@ -75,8 +84,12 @@ def training_data(self): def testing_step(self, data, target): target = np.ravel(target) prediction = self.model().predict(data) + metric_functions = { + "MEAN_SQUARE_ERROR": mean_squared_error, + "MEAN_ABSOLUTE_ERROR": mean_absolute_error, + "R2": r2_score, + } return { - "MEAN_SQUARE_ERROR": mean_squared_error(target, prediction), - "MEAN_ABSOLUTE_ERROR": mean_absolute_error(target, prediction), - "R2": r2_score(target, prediction), + name: metric_functions[name](target, prediction) + for name in self.model_args().get("metrics", self.DEFAULT_METRICS) } From f417dd43315db10595e7ff9af701262c139fcdf8 Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Thu, 29 Oct 2026 23:16:20 -0400 Subject: [PATCH 3/3] programatically request all metrics --- scripts/request_training_plan_approval.py | 8 +-- scripts/run_experiments.py | 55 +++++++++---------- .../test_request_training_plan_approval.py | 2 +- tests/scripts/test_run_experiments.py | 22 +++++--- 4 files changed, 43 insertions(+), 44 deletions(-) diff --git a/scripts/request_training_plan_approval.py b/scripts/request_training_plan_approval.py index 5d3a9ad..a68e516 100755 --- a/scripts/request_training_plan_approval.py +++ b/scripts/request_training_plan_approval.py @@ -13,7 +13,7 @@ ) # importable because the script's directory is the first entry of sys.path -from run_experiments import COMMON_TAG, TASKS, _load_training_plan_class +from run_experiments import COMMON_TAG, TASKS # Fed-BioMed-defined status of a training plan that a node doesn't know TRAINING_PLAN_NOT_REGISTERED = "Not Registered" @@ -43,9 +43,7 @@ def _request_approval(task: str) -> dict[str, str]: nodes = fds.node_ids() if fds is not None else [] if not nodes: raise ValueError(f"No node has a dataset for {task}") - training_plan = _load_training_plan_class( - config.training_plan_module, config.training_plan_class - )() + training_plan = config.training_plan_class() replies = TrainingPlanCheckJob( researcher_id=exp.researcher_id, @@ -77,7 +75,7 @@ def _request_approval(task: str) -> dict[str, str]: requests=exp.requests, nodes=nodes_not_registered, training_plan=training_plan, - description=f"{COMMON_TAG} {task}: {config.training_plan_class}", + description=f"{COMMON_TAG} {task}: {config.training_plan_class.__name__}", ).execute() for node in nodes_not_registered: not_approved[node] = STATUS_SENT diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 842fcb4..f8b73b1 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -42,10 +42,23 @@ DEFAULT_ALPHA = 0.1 +def _load_training_plan_class( + module_name: str, class_name: str +) -> type[BaseTrainingPlan]: + spec = importlib.util.spec_from_file_location( + f"training_plans.{module_name}", + DPATH_TRAINING_PLANS / f"{module_name}.py", + ) + module = importlib.util.module_from_spec(spec) + # registered so that Fed-BioMed can get the source code to send to nodes + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return getattr(module, class_name) + + @dataclass(frozen=True) class TaskConfig: - training_plan_module: str - training_plan_class: str + training_plan_class: type[BaseTrainingPlan] target_column: str # whether the model intercept is initialized to the global target mean with_target_mean: bool = False @@ -92,23 +105,19 @@ def get_fpath_out( TASKS = { "age": TaskConfig( - training_plan_module="tabular_regressor", - training_plan_class="TabularRegressorTrainingPlan", + training_plan_class=_load_training_plan_class( + "tabular_regressor", "TabularRegressorTrainingPlan" + ), target_column="nb:Age", with_target_mean=True, - extra_model_args={ - "metrics": ["MEAN_SQUARE_ERROR", "MEAN_ABSOLUTE_ERROR", "R2"] - }, secagg_clipping_range=100, ), "sex": TaskConfig( - training_plan_module="tabular_classifier", - training_plan_class="TabularClassifierTrainingPlan", + training_plan_class=_load_training_plan_class( + "tabular_classifier", "TabularClassifierTrainingPlan" + ), target_column="nb:Sex", - extra_model_args={ - "n_classes": 2, - "metrics": ["ACCURACY", "BALANCED_ACCURACY", "AUROC"], - }, + extra_model_args={"n_classes": 2}, ), } @@ -261,18 +270,6 @@ def _save_json(fpath: Path, data: dict) -> None: ) -def _load_training_plan_class(module_name: str, class_name: str) -> type: - spec = importlib.util.spec_from_file_location( - f"training_plans.{module_name}", - DPATH_TRAINING_PLANS / f"{module_name}.py", - ) - module = importlib.util.module_from_spec(spec) - # registered so that Fed-BioMed can get the source code to send to nodes - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return getattr(module, class_name) - - def _get_node_data(task: str) -> dict[str, dict]: """Get the dataset metadata of each node that has data for a task.""" fds = Experiment(tags=[COMMON_TAG, task]).training_data() @@ -345,6 +342,8 @@ def _build_model_args( "feature_std": std, "target_column": config.target_column, "alpha": alpha, + # all the metrics supported by the training plan + "metrics": list(config.training_plan_class.METRICS), **config.extra_model_args, } if config.with_target_mean: @@ -511,11 +510,7 @@ def _run_task( print(f"[{task}] Using {len(columns)} features") model_args = _build_model_args(task, columns, mean, std, alpha, target_mean) exp.set_model_args(model_args) - exp.set_training_plan_class( - _load_training_plan_class( - config.training_plan_module, config.training_plan_class - ) - ) + exp.set_training_plan_class(config.training_plan_class) # saved before training so that they are available even if training fails json_data = { diff --git a/tests/scripts/test_request_training_plan_approval.py b/tests/scripts/test_request_training_plan_approval.py index 7344a2e..42672b2 100644 --- a/tests/scripts/test_request_training_plan_approval.py +++ b/tests/scripts/test_request_training_plan_approval.py @@ -93,7 +93,7 @@ def test_approved( assert kwargs["nodes"] == NODES training_plan = kwargs["training_plan"] assert isinstance(training_plan, BaseTrainingPlan) - assert type(training_plan).__name__ == module.TASKS[task].training_plan_class + assert type(training_plan) is module.TASKS[task].training_plan_class mock_approve_job.assert_not_called() def test_not_approved( diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 119b16e..d7fdbce 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -143,12 +143,9 @@ def fpath_csv(tmp_path: Path) -> Path: class TestLoadTrainingPlanClass: @pytest.mark.parametrize("task", ["age", "sex"]) def test_source(self, module: ModuleType, task: str): - config = module.TASKS[task] - plan_cls = module._load_training_plan_class( - config.training_plan_module, config.training_plan_class - ) + plan_cls = module.TASKS[task].training_plan_class assert issubclass(plan_cls, BaseTrainingPlan) - assert f"class {config.training_plan_class}" in plan_cls().source() + assert f"class {plan_cls.__name__}" in plan_cls().source() @pytest.fixture @@ -418,9 +415,7 @@ def test_training_data( task, columns, mean, std, 0.1, target_mean ) - plan = module._load_training_plan_class( - config.training_plan_module, config.training_plan_class - )() + plan = config.training_plan_class() plan.post_init(model_args, TrainingArgs({}, only_required=False)) data_manager = plan.training_data() data_manager.load(tp_type=plan.type()) @@ -441,6 +436,17 @@ def test_target_mean(self, module: ModuleType): model_args = module._build_model_args("age", ["a"], [0.0], [1.0], 0.1, 30.0) assert model_args["target_mean"] == 30.0 + @pytest.mark.parametrize( + "task,expected_metrics", + [ + ("age", ["MEAN_SQUARE_ERROR", "MEAN_ABSOLUTE_ERROR", "R2"]), + ("sex", ["ACCURACY", "BALANCED_ACCURACY", "AUROC"]), + ], + ) + def test_metrics(self, module: ModuleType, task: str, expected_metrics: list): + model_args = module._build_model_args(task, ["a"], [0.0], [1.0], 0.1, 30.0) + assert model_args["metrics"] == expected_metrics + class TestBuildOptimizerArgs: @pytest.mark.parametrize(