Hotfix fixed batch_size

README.rst

@@ -196,6 +196,10 @@ Changelogs
 HSF
 ---
 
+**Version 1.0.1**
+
+* Fix batch size issue
+
 **Version 1.0.0**
 
 * Added Uncertainty Maps for post-hoc analysis of segmentation results,

docs/about/release-notes.md

@@ -16,6 +16,10 @@ Current maintainers:
 
 ## HSF
 
+### Version 1.0.1 (2021-12-07)
+
+* Fixed batch size issue
+
 ### Version 1.0.0 (2021-11-12)
 
 * Added Uncertainty Maps for post-hoc analysis of segmentation results,

docs/index.md

@@ -5,7 +5,7 @@
     <br>
     <font size="+2"><b>Hippocampal</b> <i>Segmentation</i> Factory</font>
     <br>
-    <b>Current HSF version:</b> 1.0.0<br>
+    <b>Current HSF version:</b> 1.0.1<br>
     <b>Built-in Models version:</b> 2.0.0<br>
     <b>Models in the Hub:</b> 4
 </p>

hsf/__init__.py

@@ -1 +1 @@
-__version__ = '1.0.0'
+__version__ = '1.0.1'

hsf/engines.py

@@ -1,4 +1,5 @@
 from pathlib import Path, PosixPath
+from typing import Generator
 
 import deepsparse
 import onnxruntime as ort
@@ -41,7 +42,7 @@ def print_deepsparse_support():
 
 
 def get_inference_engines(models_path: PosixPath, engine_name: str,
-                          engine_settings: DictConfig) -> list:
+                          engine_settings: DictConfig) -> Generator:
     """
     Returns Inference Engines.
 
@@ -56,11 +57,10 @@ def get_inference_engines(models_path: PosixPath, engine_name: str,
     p = Path(models_path).expanduser()
     models = list(p.glob("*.onnx"))
 
-    return [
-        InferenceEngine(engine_name=engine_name,
-                        engine_settings=engine_settings,
-                        model=model) for model in models
-    ]
+    for model in models:
+        yield InferenceEngine(engine_name=engine_name,
+                              engine_settings=engine_settings,
+                              model=model)
 
 
 class InferenceEngine:

hsf/factory.py

@@ -70,13 +70,13 @@ def predict(mri: PosixPath, engines: list, cfg: DictConfig) -> tuple:
     subject = mri_to_subject(mri)
 
     pprint(f"{PREFIX} Starting segmentation...")
-    return segment(
-        subject=subject,
-        augmentation_cfg=cfg.augmentation,
-        segmentation_cfg=cfg.segmentation.segmentation,
-        engines=engines,
-        ca_mode=str(cfg.segmentation.ca_mode),
-    )
+    return segment(subject=subject,
+                   augmentation_cfg=cfg.augmentation,
+                   segmentation_cfg=cfg.segmentation.segmentation,
+                   n_engines=len(cfg.segmentation.models),
+                   engines=engines,
+                   ca_mode=str(cfg.segmentation.ca_mode),
+                   batch_size=cfg.hardware.engine_settings.batch_size)
 
 
 def compute_uncertainty(mri: PosixPath, soft_pred: torch.Tensor) -> None:
@@ -132,11 +132,10 @@ def save(mri: PosixPath, hippocampus: PosixPath, hard_pred: torch.Tensor,
 def main(cfg: DictConfig) -> None:
     fetch_models(cfg.segmentation.models_path, cfg.segmentation.models)
 
-    engines = get_inference_engines(
-        cfg.segmentation.models_path,
-        engine_name=cfg.hardware.engine,
-        engine_settings=cfg.hardware.engine_settings)
-    pprint(f"{PREFIX} Successfully loaded segmentation models in memory.")
+    if cfg.hardware.engine == "deepsparse":
+        tta = cfg.segmentation.segmentation.test_time_num_aug
+        bs = cfg.hardware.engine_settings.batch_size
+        assert tta % bs == 0, "test_time_num_aug must be a multiple of batch_size for deepsparse"
 
     mris = load_from_config(cfg.files.path, cfg.files.pattern)
 
@@ -157,6 +156,11 @@ def main(cfg: DictConfig) -> None:
         locator, orientation, hippocampi = get_lr_hippocampi(mri, cfg)
 
         for j, hippocampus in enumerate(hippocampi):
+            engines = get_inference_engines(
+                cfg.segmentation.models_path,
+                engine_name=cfg.hardware.engine,
+                engine_settings=cfg.hardware.engine_settings)
+
             hippocampus = Path(hippocampus)
 
             pprint(f"{PREFIX} Subject {i+1}/{N}, side {j+1}/2")

hsf/segment.py

@@ -96,14 +96,15 @@ def predict(mris: list,
         aug.add_image(lm_temp, 'label')
         aug.label.set_data(lab)
         back = aug.apply_inverse_transform(warn=True)
-        results.extend(back.label.data)
+        results.append(back.label.data)
 
-    return torch.stack(results, dim=0)
+    return results
 
 
 def segment(subject: tio.Subject,
             augmentation_cfg: DictConfig,
             segmentation_cfg: DictConfig,
+            n_engines: int,
             engines: list,
             ca_mode: str = "1/2/3",
             batch_size: int = 1) -> tuple:
@@ -129,16 +130,15 @@ def segment(subject: tio.Subject,
     ]
 
     results = []
-    for sub in track(
-            batched_subjects,
-            description=
-            f"Segmenting (TTA: {len(subjects)} | {len(engines)} MODELS)..."):
-        engines_predictions = [
-            predict(sub, engine, ca_mode) for engine in engines
-        ]
-
-        results.extend(engines_predictions)
-
+    n = 0
+    for engine in engines:
+        n += 1
+        for sub in track(
+                batched_subjects,
+                description=
+                f"Segmenting (TTA: {len(subjects)} | MODEL {n}/{n_engines})..."
+        ):
+            results.extend(predict(sub, engine, ca_mode))
     soft_predictions = torch.stack(results, dim=0)
     hard_prediction = soft_predictions.argmax(dim=1).long().mode(dim=0).values
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "HSF"
-version = "1.0.0"
+version = "1.0.1"
 description = "A simple yet exhaustive segmentation tool of the Hippocampal Subfields in T1w and T2w MRIs."
 authors = ["Clément POIRET <poiret.clement@outlook.fr>"]
 license = "MIT"

tests/test_hsf.py

@@ -15,7 +15,7 @@
 
 
 def test_version():
-    assert __version__ == '1.0.0'
+    assert __version__ == '1.0.1'
 
 
 # SETUP FIXTURES