[ENH] Add custom metrics to training plan

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

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 = {

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(

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(

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

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}