Working Apple Silicon MPS workaround for DentalSegmentator / SlicerNNUnet under Rosetta

Development
1

I tested a proof of concept that allows DentalSegmentator to use Apple Silicon GPU acceleration through MPS, while the main 3D Slicer application continues running as x86_64 under Rosetta.

Environment

  • MacBook Pro with Apple M3 Pro
  • 3D Slicer 5.12.3 stable, x86_64
  • DentalSegmentator with SlicerNNUnet
  • nnUNetv2 2.8.1
  • Slicer-bundled PyTorch 2.2.2

Problem

Slicer reports that MPS is built and available, and DentalSegmentator accepts mps as its device. However, inference fails with:

RuntimeError: Conv3D is not supported on MPS

The problem appears to be the older x86_64 PyTorch build bundled with Slicer, rather than DentalSegmentator or nnUNet itself.

Proof of concept

I created a separate native arm64 Python environment containing:

  • PyTorch 2.13.0
  • torchvision 0.28.0
  • nnUNetv2 2.8.1

A Conv3D smoke test succeeds on mps:0 in this environment.

SlicerNNUnet was then pointed to a small wrapper:

#!/bin/sh
unset PYTHONHOME PYTHONPATH
export PYTORCH_ENABLE_MPS_FALLBACK=1
exec "/path/to/native-arm64-venv/bin/nnUNetv2_predict" "$@"

Clearing PYTHONHOME and PYTHONPATH is necessary because otherwise the external interpreter imports nnUNet and dependencies from Slicer’s x86_64 Python environment.

For the proof of concept, I overrode SegmentationLogic._findUNetPredictPath from .slicerrc.py so that it returned the wrapper path.

Result

  • Complete DentalSegmentator run from the normal Slicer interface
  • Device: MPS
  • 252/252 inference patches completed
  • Approximately 1.28 seconds per patch
  • Approximately 7 minutes total
  • CPU execution on the same system was approximately 9 seconds per patch, with an estimated total of 39 minutes
  • The resulting segmentation was successfully loaded back into Slicer

This was tested on one Apple M3 Pro and one dental CT volume, so these numbers should not be treated as a formal benchmark. No image or patient data are being shared.

Possible upstream improvement

Would the maintainers consider adding a supported advanced setting or environment variable for overriding the nnUNetv2_predict executable path?

This would allow SlicerNNUnet-based extensions to delegate inference to a native arm64 environment without requiring a full native Apple Silicon build of Slicer. The current _findUNetPredictPath override is useful as a proof of concept but relies on an internal method.

Relevant references:

I can provide the complete installation commands and help test a cleaner implementation if the approach is considered useful.