Fixes ensembling with sigmoid (#6588)

Proper ensembling when trained with a sigmoid (overlapping classes) in Auto3dseg. --------- Signed-off-by: myron <[email protected]>
Project-MONAI · Jun 8, 2023 · e6ec945 · e6ec945
1 parent 52b3ed2
commit e6ec945
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/monai/apps/auto3dseg/auto_runner.py b/monai/apps/auto3dseg/auto_runner.py
@@ -292,6 +292,9 @@ def __init__(
         self.search_space: dict[str, dict[str, Any]] = {}
         self.hpo_tasks = 0
 
+        if "sigmoid" not in self.kwargs and "sigmoid" in self.data_src_cfg:
+            self.kwargs["sigmoid"] = self.data_src_cfg["sigmoid"]
+
     def read_cache(self):
         """
         Check if the intermediate result is cached after each step in the current working directory

diff --git a/monai/apps/auto3dseg/ensemble_builder.py b/monai/apps/auto3dseg/ensemble_builder.py
@@ -586,14 +586,10 @@ def ensemble(self):
             if len(infer_files) == 0:
                 logger.info("No testing files for inference is provided. Ensembler ending.")
                 return
-            infer_files = (
-                partition_dataset(data=infer_files, shuffle=False, num_partitions=len(infer_files))[self.rank]
-                if self.rank < len(infer_files)
-                else []
-            )
+            infer_files = [infer_files[self.rank]] if self.rank < len(infer_files) else []
         else:
             infer_files = partition_dataset(
-                data=infer_files, shuffle=False, num_partitions=self.world_size, even_divisible=True
+                data=infer_files, shuffle=False, num_partitions=self.world_size, even_divisible=False
             )[self.rank]
 
         # TO DO: Add some function in ensembler for infer_files update?