Plan CPU-only sliding window

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -943,6 +943,9 @@ class SlidingWindowConfig:
     padding_mode: str = "constant"  # Padding mode at volume boundaries
     pad_size: Optional[List[int]] = None  # Padding size for context (e.g., [16, 32, 32])
     save_channels: Optional[List[int]] = None  # Channel indices to save (None = all channels)
+    keep_input_on_cpu: bool = False  # Keep full input volume on CPU; MONAI moves windows to sw_device
+    sw_device: Optional[str] = None  # MONAI window compute device (e.g., "cuda")
+    output_device: Optional[str] = None  # MONAI stitched output device (e.g., "cpu")
 
 
 @dataclass

connectomics/inference/sliding.py

@@ -10,6 +10,7 @@
 from typing import Optional, Tuple, Union
 import warnings
 
+import torch
 from monai.inferers import SlidingWindowInferer
 
 
@@ -97,6 +98,28 @@ def build_sliding_inferer(cfg):
     mode = getattr(sliding_cfg, "blending", "gaussian") if sliding_cfg else "gaussian"
     sigma_scale = float(getattr(sliding_cfg, "sigma_scale", 0.125)) if sliding_cfg else 0.125
     padding_mode = getattr(sliding_cfg, "padding_mode", "constant") if sliding_cfg else "constant"
+    keep_input_on_cpu = bool(getattr(sliding_cfg, "keep_input_on_cpu", False)) if sliding_cfg else False
+    sw_device = getattr(sliding_cfg, "sw_device", None) if sliding_cfg else None
+    output_device = getattr(sliding_cfg, "output_device", None) if sliding_cfg else None
+
+    if isinstance(sw_device, str) and sw_device.lower() in {"", "none", "null"}:
+        sw_device = None
+    if isinstance(output_device, str) and output_device.lower() in {"", "none", "null"}:
+        output_device = None
+
+    if keep_input_on_cpu:
+        # MONAI defaults both devices to inputs.device; when inputs stay on CPU,
+        # explicitly route windows to CUDA and stitch on CPU unless overridden.
+        if sw_device is None and torch.cuda.is_available():
+            sw_device = "cuda"
+        if output_device is None:
+            output_device = "cpu"
+        if sw_device is None:
+            warnings.warn(
+                "inference.sliding_window.keep_input_on_cpu=True but no sw_device was set "
+                "and CUDA is unavailable. Sliding-window inference will run on CPU.",
+                UserWarning,
+            )
 
     inferer = SlidingWindowInferer(
         roi_size=roi_size,
@@ -105,13 +128,16 @@ def build_sliding_inferer(cfg):
         mode=mode,
         sigma_scale=sigma_scale,
         padding_mode=padding_mode,
+        sw_device=sw_device,
+        device=output_device,
         progress=True,
     )
 
     print(
         "  Sliding-window inference configured: "
         f"roi_size={roi_size}, overlap={overlap}, sw_batch={sw_batch_size}, "
-        f"mode={mode}, sigma_scale={sigma_scale}, padding={padding_mode}"
+        f"mode={mode}, sigma_scale={sigma_scale}, padding={padding_mode}, "
+        f"keep_input_on_cpu={keep_input_on_cpu}, sw_device={sw_device}, output_device={output_device}"
     )
 
     return inferer

connectomics/inference/tta.py

@@ -168,6 +168,18 @@ def _run_network(self, images: torch.Tensor) -> torch.Tensor:
         with torch.no_grad():
             if self.sliding_inferer is not None:
                 return self.sliding_inferer(inputs=images, network=self._sliding_window_predict)
+            keep_input_on_cpu = bool(
+                getattr(
+                    getattr(getattr(self.cfg, "inference", None), "sliding_window", None),
+                    "keep_input_on_cpu",
+                    False,
+                )
+            )
+            if keep_input_on_cpu and images.device.type == "cpu":
+                raise RuntimeError(
+                    "inference.sliding_window.keep_input_on_cpu=True requires sliding-window "
+                    "inference to be enabled (set inference.sliding_window.window_size or model output size)."
+                )
             return self._sliding_window_predict(images)
 
     def _sliding_window_predict(self, inputs: torch.Tensor) -> torch.Tensor:
@@ -399,6 +411,8 @@ def predict(self, images: torch.Tensor, mask: Optional[torch.Tensor] = None) ->
             else False
         )
         if apply_mask and mask is not None:
+            if mask.device != ensemble_result.device:
+                mask = mask.to(device=ensemble_result.device, non_blocking=True)
             if mask.shape != ensemble_result.shape:
                 if not (mask.shape[1] == 1 and mask.shape[0] == ensemble_result.shape[0]):
                     warnings.warn(

connectomics/training/lit/model.py

@@ -396,7 +396,7 @@ def _save_metrics_to_file(self, metrics_dict: Dict[str, Any]):
     def _compute_test_metrics(self, decoded_predictions: np.ndarray, labels: torch.Tensor, volume_name: str = None):
         """Update configured torchmetrics using decoded predictions and print per-volume metrics."""
         pred_tensor = torch.from_numpy(decoded_predictions).float().to(self.device)
-        labels_tensor = labels.float()
+        labels_tensor = labels.float().to(pred_tensor.device)
 
         # Remove batch and channel dimensions
         pred_tensor = pred_tensor.squeeze()
@@ -770,13 +770,57 @@ def on_test_start(self):
             # Keep in training mode (e.g., for Monte Carlo Dropout uncertainty estimation)
             self.train()
 
+        sliding_cfg = getattr(getattr(self.cfg, "inference", None), "sliding_window", None)
+        if bool(getattr(sliding_cfg, "keep_input_on_cpu", False)):
+            print(
+                "  Sliding-window CPU input mode enabled: keeping test image tensors on CPU "
+                "and letting MONAI move window batches to the configured sw_device."
+            )
+
     def on_test_end(self):
         """Called at the end of testing."""
         # Note: Metrics are logged in test_step with on_epoch=True,
         # which automatically computes and logs them at the end of testing.
         # Logging here causes a Lightning warning since self.log() is not allowed in on_test_end.
         pass
 
+    def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any:
+        """Keep large test/predict input volumes on CPU for MONAI sliding-window inference."""
+        sliding_cfg = getattr(getattr(self.cfg, "inference", None), "sliding_window", None)
+        keep_input_on_cpu = bool(getattr(sliding_cfg, "keep_input_on_cpu", False))
+
+        trainer = getattr(self, "_trainer", None)
+        is_test_or_predict = bool(
+            getattr(trainer, "testing", False) or getattr(trainer, "predicting", False)
+        )
+
+        preserve_cpu_input = keep_input_on_cpu and is_test_or_predict and isinstance(batch, dict)
+        cpu_image = None
+        cpu_label = None
+        cpu_mask = None
+        if preserve_cpu_input:
+            image = batch.get("image")
+            label = batch.get("label")
+            mask = batch.get("mask")
+            if torch.is_tensor(image):
+                cpu_image = image
+            if torch.is_tensor(label):
+                cpu_label = label
+            if torch.is_tensor(mask):
+                cpu_mask = mask
+
+        moved_batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
+
+        if preserve_cpu_input and isinstance(moved_batch, dict):
+            if cpu_image is not None:
+                moved_batch["image"] = cpu_image
+            if cpu_label is not None:
+                moved_batch["label"] = cpu_label
+            if cpu_mask is not None:
+                moved_batch["mask"] = cpu_mask
+
+        return moved_batch
+
     def test_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STEP_OUTPUT:
         """
         Test step with optional sliding-window inference and metrics computation.
@@ -885,6 +929,7 @@ def test_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STEP_OUTP
         print(f"INFERENCE PLAN: {volume_name}")
         print(f"{'='*70}")
         print(f"Input shape:       {tuple(images.shape)}")
+        print(f"Input device:      {images.device}")
 
         # Extract sliding window parameters if available
         if hasattr(self.cfg, 'inference') and hasattr(self.cfg.inference, 'sliding_window'):

tutorials/neuron_nisb_base_40nm.yaml

@@ -170,11 +170,14 @@ test:
 inference:
   sliding_window:
     window_size: [128, 128, 128]
-    sw_batch_size: 4
+    sw_batch_size: 1
     overlap: 0.5
     blending: gaussian
     sigma_scale: 0.25
     padding_mode: replicate
+    keep_input_on_cpu: true
+    sw_device: cuda
+    output_device: cpu
 
   test_time_augmentation:
     enabled: false