Repository: Project-MONAI/MONAI

Issue:
auto3dseg flexibility to override inference parameters for each model in the ensemble
**Is your feature request related to a problem? Please describe.**
it seems in the ensemble logic, `pred_param` should be given in a way that's compatible with all the models participating in the inference.

https://github.com/Project-MONAI/MONAI/blob/eb6bbe8da518f49abd001ca5865672933ab1e9a2/monai/apps/auto3dseg/ensemble_builder.py#L116-L118

would be great to revise the following logic to support model-specific params as well.

https://github.com/Project-MONAI/MONAI/blob/eb6bbe8da518f49abd001ca5865672933ab1e9a2/monai/apps/auto3dseg/ensemble_builder.py#L158-L161


Discussion:

Resolving patch:
diff --git a/monai/apps/auto3dseg/__init__.py b/monai/apps/auto3dseg/__init__.py
--- a/monai/apps/auto3dseg/__init__.py
+++ b/monai/apps/auto3dseg/__init__.py
@@ -22,4 +22,4 @@
     EnsembleRunner,
 )
 from .hpo_gen import NNIGen, OptunaGen
-from .utils import export_bundle_algo_history, import_bundle_algo_history
+from .utils import export_bundle_algo_history, get_name_from_algo_id, import_bundle_algo_history
diff --git a/monai/apps/auto3dseg/ensemble_builder.py b/monai/apps/auto3dseg/ensemble_builder.py
--- a/monai/apps/auto3dseg/ensemble_builder.py
+++ b/monai/apps/auto3dseg/ensemble_builder.py
@@ -24,7 +24,7 @@
 import torch.distributed as dist
 
 from monai.apps.auto3dseg.bundle_gen import BundleAlgo
-from monai.apps.auto3dseg.utils import import_bundle_algo_history
+from monai.apps.auto3dseg.utils import get_name_from_algo_id, import_bundle_algo_history
 from monai.apps.utils import get_logger
 from monai.auto3dseg import concat_val_to_np
 from monai.auto3dseg.utils import datafold_read
@@ -127,6 +127,25 @@ def ensemble_pred(self, preds, sigmoid=False):
             else:
                 return VoteEnsemble(num_classes=preds[0].shape[0])(classes)
 
+    def _apply_algo_specific_param(self, algo_spec_param: dict, param: dict, algo_name: str) -> dict:
+        """
+        Apply the model-specific params to the prediction params based on the name of the Algo.
+
+        Args:
+            algo_spec_param: a dict that has structure of {"<name of algo>": "<pred_params for that algo>"}.
+            param: the prediction params to override.
+            algo_name: name of the Algo
+
+        Returns:
+            param after being updated with the model-specific param
+        """
+        _param_to_override = deepcopy(algo_spec_param)
+        _param = deepcopy(param)
+        for k, v in _param_to_override.items():
+            if k.lower() == algo_name.lower():
+                _param.update(v)
+        return _param
+
     def __call__(self, pred_param: dict | None = None) -> list:
         """
         Use the ensembled model to predict result.
@@ -146,6 +165,8 @@ def __call__(self, pred_param: dict | None = None) -> list:
                       Example: `{"_target_": "SaveImage", "output_dir": "./"}`
                     - ``"sigmoid"``: use the sigmoid function (e.g. x > 0.5) to convert the prediction probability map
                       to the label class prediction, otherwise argmax(x) is used.
+                    - ``"algo_spec_params"``: a dictionary to add pred_params that are specific to a model.
+                      The dict has a format of {"<name of algo>": "<pred_params for that algo>"}.
 
                 The parameters in the second group is defined in the ``config`` of each Algo templates. Please check:
                 https://github.com/Project-MONAI/research-contributions/tree/main/auto3dseg/algorithm_templates
@@ -172,6 +193,8 @@ def __call__(self, pred_param: dict | None = None) -> list:
         if "image_save_func" in param:
             img_saver = ConfigParser(param["image_save_func"]).get_parsed_content()
 
+        algo_spec_params = param.pop("algo_spec_params", {})
+
         outputs = []
         for _, file in (
             enumerate(tqdm(files, desc="Ensembling (rank 0)..."))
@@ -180,8 +203,10 @@ def __call__(self, pred_param: dict | None = None) -> list:
         ):
             preds = []
             for algo in self.algo_ensemble:
+                infer_algo_name = get_name_from_algo_id(algo[AlgoKeys.ID])
                 infer_instance = algo[AlgoKeys.ALGO]
-                pred = infer_instance.predict(predict_files=[file], predict_params=param)
+                _param = self._apply_algo_specific_param(algo_spec_params, param, infer_algo_name)
+                pred = infer_instance.predict(predict_files=[file], predict_params=_param)
                 preds.append(pred[0])
             if "image_save_func" in param:
                 try:
diff --git a/monai/apps/auto3dseg/utils.py b/monai/apps/auto3dseg/utils.py
--- a/monai/apps/auto3dseg/utils.py
+++ b/monai/apps/auto3dseg/utils.py
@@ -17,6 +17,8 @@
 from monai.auto3dseg import algo_from_pickle, algo_to_pickle
 from monai.utils.enums import AlgoKeys
 
+__all__ = ["import_bundle_algo_history", "export_bundle_algo_history", "get_name_from_algo_id"]
+
 
 def import_bundle_algo_history(
     output_folder: str = ".", template_path: str | None = None, only_trained: bool = True
@@ -73,3 +75,16 @@ def export_bundle_algo_history(history: list[dict[str, BundleAlgo]]) -> None:
     for algo_dict in history:
         algo = algo_dict[AlgoKeys.ALGO]
         algo_to_pickle(algo, template_path=algo.template_path)
+
+
+def get_name_from_algo_id(id: str) -> str:
+    """
+    Get the name of Algo from the identifier of the Algo.
+
+    Args:
+        id: identifier which follows a convention of "name_fold_other".
+
+    Returns:
+        name of the Algo.
+    """
+    return id.split("_")[0]


Test patch:
diff --git a/tests/test_auto3dseg_ensemble.py b/tests/test_auto3dseg_ensemble.py
--- a/tests/test_auto3dseg_ensemble.py
+++ b/tests/test_auto3dseg_ensemble.py
@@ -72,7 +72,12 @@
     else {}
 )
 
-pred_param = {"files_slices": slice(0, 1), "mode": "mean", "sigmoid": True}
+pred_param = {
+    "files_slices": slice(0, 1),
+    "mode": "mean",
+    "sigmoid": True,
+    "algo_spec_params": {"segresnet": {"network#init_filters": 8}, "swinunetr": {"network#feature_size": 12}},
+}
 
 
 def create_sim_data(dataroot, sim_datalist, sim_dim, **kwargs):
@@ -171,11 +176,6 @@ def test_ensemble(self) -> None:
         builder = AlgoEnsembleBuilder(history, data_src_cfg_name=self.data_src_cfg_name)
         builder.set_ensemble_method(AlgoEnsembleBestN(n_best=1))
         ensemble = builder.get_ensemble()
-        name = ensemble.get_algo_ensemble()[0][AlgoKeys.ID]
-        if name.startswith("segresnet"):
-            pred_param["network#init_filters"] = 8
-        elif name.startswith("swinunetr"):
-            pred_param["network#feature_size"] = 12
         preds = ensemble(pred_param)
         self.assertTupleEqual(preds[0].shape, (2, 24, 24, 24))