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 3b6f32a..9127922 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -17,7 +17,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 @@ -45,12 +44,24 @@ DEFAULT_DATE_TAG = date.today().strftime("%Y_%m_%d") +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 - 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,18 +108,18 @@ 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", - test_metric=MetricTypes.MEAN_SQUARE_ERROR, with_target_mean=True, 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", - test_metric=MetricTypes.ACCURACY, extra_model_args={"n_classes": 2}, ), } @@ -270,18 +281,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() @@ -354,6 +353,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: @@ -480,7 +481,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( @@ -521,11 +521,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 02fe975..768fa6e 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 @@ -420,9 +417,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()) @@ -443,6 +438,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( @@ -579,7 +585,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..ebeba5c 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,79 @@ 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 | {"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]]) + 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") + @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 | {"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]])) + 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( + 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..00d8a25 100644 --- a/tests/training_plans/test_tabular_regressor.py +++ b/tests/training_plans/test_tabular_regressor.py @@ -124,3 +124,63 @@ 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 | {"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])} + ) + # 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) + + 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 fe39afa..6d6ddfe 100644 --- a/training_plans/tabular_classifier.py +++ b/training_plans/tabular_classifier.py @@ -7,11 +7,16 @@ 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): """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", @@ -20,6 +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, roc_auc_score", ] def init_optimizer(self, optimizer_args): @@ -37,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 @@ -63,3 +77,18 @@ 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) + metrics = { + "ACCURACY": accuracy_score(target, prediction), + "BALANCED_ACCURACY": balanced_accuracy_score(target, prediction), + "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 751300a..9058ef7 100644 --- a/training_plans/tabular_regressor.py +++ b/training_plans/tabular_regressor.py @@ -7,11 +7,16 @@ 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): """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", @@ -20,6 +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 mean_absolute_error, mean_squared_error, r2_score", ] def init_optimizer(self, optimizer_args): @@ -39,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"])]) @@ -46,7 +57,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 +80,16 @@ 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) + metric_functions = { + "MEAN_SQUARE_ERROR": mean_squared_error, + "MEAN_ABSOLUTE_ERROR": mean_absolute_error, + "R2": r2_score, + } + return { + name: metric_functions[name](target, prediction) + for name in self.model_args().get("metrics", self.DEFAULT_METRICS) + }