Fixes for FastSurfer handles running out of GPU memory

[santiestrada32]

Summary

A batch-size × voxel-size sweep on an RTX 4080 SUPER (16 GB) showed three problems with how FastSurfer handles running out of GPU memory:

1. Inconsistent flag name. run_fastsurfer.sh (and so the Docker/Singularity images) only accepted --batch. The Python modules, and the OOM hint itself, use --batch_size. Following the hint gave ERROR: Flag '--batch_size' unrecognized.
2. An OOM exited with code 0. FastSurferCNN/run_prediction.py caught the CUDA OOM error, printed the hint and then returned 0. run_fastsurfer.sh carried on as if the step had passed, even though no segmentation was written. All 13 runs in the sweep that ran out of memory exited with code 0.
3. Advice that doesn't apply. The hint always suggested a smaller batch size, even at the default of 1. At 0.4–0.5 mm, runs fail even at batch size 1, so a larger --vox_size is the advice that actually helps there.

Changes

• run_fastsurfer.sh, recon_surf/long_prepare_template.sh: accept --batch_size <int>, and update the help text to match. --batch still works but prints a deprecation warning, following the existing pattern (e.g. --run_viewagg_on). brun_fastsurfer.sh/srun_fastsurfer.sh keep their own --batch "<i>/<n>", which means task splitting. That flag used to clash with the batch size; --batch_size can now be passed through brun.
• FastSurferCNN/run_prediction.py: on a CUDA OOM, return "Segmentation failed: insufficient GPU memory.", so the script exits 1 and run_fastsurfer.sh stops with ERROR: FastSurfer asegdkt segmentation failed.
• FastSurferCNN/utils/common.py: handle_cuda_memory_exception takes an optional batch_size.
  • batch size > 1, or not passed in: suggest a smaller --batch_size, a larger --vox_size, and running view aggregation on the CPU (--viewagg_device 'cpu').
  • batch size == 1: suggest only a larger --vox_size and --viewagg_device 'cpu'.
• FastSurferCNN/run_prediction.py, CerebNet/inference.py: pass the batch size to the handler.

Testing

Tested with the v2.6.0 image (fastsurfer-cu126-v260.sif), with this branch mounted over /fastsurfer, on a 0.8 mm T1. All runs used --seg_only --no_cereb --no_hypothal --no_cc --no_biasfield.

Run	Exit code	Result
--vox_size 0.5 --batch_size 1	1	OOM, hint without the batch-size advice, asegdkt segmentation failed
--vox_size 0.5 --batch_size 4	1	OOM, hint with the --batch_size advice, asegdkt segmentation failed
--vox_size 1.0 --batch_size 1	0	segmentation written
--vox_size 1.0 --batch 1	0	segmentation written, deprecation warning printed

The two shell scripts pass bash -n. Calling handle_cuda_memory_exception directly with batch_size=1, 8 and None logs the expected wording for each.

[Copilot]

🟢 Approval recommended

The changes consistently address the flag mismatch, failure status, and misleading OOM advice without identified regressions.

0 open findings

What changed in this PR

Improves GPU OOM handling and standardizes the inference batch-size flag.

Changes:

• Adds --batch_size, retaining deprecated --batch compatibility.
• Makes FastSurferCNN OOM failures exit nonzero.
• Tailors OOM guidance to the active batch size.

File	Description
run_fastsurfer.sh	Adds the standardized flag and deprecation warning.
recon_surf/​long_prepare_template.sh	Supports and forwards the standardized flag.
FastSurferCNN/​utils/​common.py	Provides batch-aware OOM guidance.
FastSurferCNN/​run_prediction.py	Returns an error after CUDA OOM.
CerebNet/​inference.py	Supplies batch size to the OOM handler.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.