Running nnUnet inference on Apple GPU

[valosekj]

I tried to install nnU-Net locally on my new MacBook using:

model-spinal-rootlets/packaging_lumbar_rootlets/README.md

Line 44 in 72c1842

pip install -r packaging_lumbar_rootlets/requirements.txt

But I got this error when running `run_inference_single_subject.py`

 python ~/code/model-spinal-rootlets/packaging_lumbar_rootlets/run_inference_single_subject.py -h
Traceback (most recent call last):
  File "/Users/valosek/code/model-spinal-rootlets/packaging_ventral_rootlets/run_inference_single_subject.py", line 31, in <module>
    from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
  File "/Users/valosek/miniconda3/envs/nnunet/lib/python3.9/site-packages/nnunetv2/inference/predict_from_raw_data.py", line 22, in <module>
    from nnunetv2.inference.data_iterators import PreprocessAdapterFromNpy, preprocessing_iterator_fromfiles, \
  File "/Users/valosek/miniconda3/envs/nnunet/lib/python3.9/site-packages/nnunetv2/inference/data_iterators.py", line 12, in <module>
    from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor
  File "/Users/valosek/miniconda3/envs/nnunet/lib/python3.9/site-packages/nnunetv2/preprocessing/preprocessors/default_preprocessor.py", line 25, in <module>
    from nnunetv2.preprocessing.cropping.cropping import crop_to_nonzero
  File "/Users/valosek/miniconda3/envs/nnunet/lib/python3.9/site-packages/nnunetv2/preprocessing/cropping/cropping.py", line 5, in <module>
    from acvl_utils.cropping_and_padding.bounding_boxes import get_bbox_from_mask, crop_to_bbox, bounding_box_to_slice
  File "/Users/valosek/miniconda3/envs/nnunet/lib/python3.9/site-packages/acvl_utils/cropping_and_padding/bounding_boxes.py", line 5, in <module>
    import blosc2
  File "/Users/valosek/miniconda3/envs/nnunet/lib/python3.9/site-packages/blosc2/__init__.py", line 86, in <module>
    from .blosc2_ext import (
  File "blosc2_ext.pyx", line 1, in init blosc2.blosc2_ext
ValueError: numpy.dtype size changed, may indicate binary incompatibility. Expected 96 from C header, got 88 from PyObject

So, I tried `venv_sct`, and it worked!

source $SCT_DIR/python/etc/profile.d/conda.sh
conda activate venv_sct
python ~/code/model-spinal-rootlets/packaging_lumbar_rootlets/run_inference_single_subject.py -h

But since it took several minutes to run inference on a single image using CPU, I checked the status of accelaring nnunet on Apple's GPUs and found this super relevant PR; continued in the next comment.

[valosekj]

nnunet inference on Apple's GPUs

Install pytorch supporting 3D convolution on MPS

git clone --depth 1 --branch convtranspose_mps_remove_check https://github.com/mlaves/pytorch.git
cd pytorch
git submodule update --init --recursive

python3 -m venv venv
source ./venv/bin/activate
pip install -r requirements.txt
python3 setup.py develop

(source)

Install nnunet

pip install nnunetv2

Run your model

Run your model using run_inference_single_subject.py (with device=torch.device('mps')):

python run_inference_single_subject.py -i sub-01_T2w.nii.gz -path-model <PATH_TO_YOUR_MODEL>/nnUNetTrainerDA5__nnUNetPlans__3d_fullres -fold 0 -o sub-01_T2w_seg.nii.gz

[joshuacwnewton]

git clone --depth 1 --branch convtranspose_mps_remove_check https://github.com/mlaves/pytorch.git
cd pytorch
git submodule update --init --recursive

Just to make sure I'm understanding correctly: The need for using this fork is (hopefully) temporary, since the above branch convtranspose_mps_remove_check directly corresponds to this open PR that may very well be merged into the mainline torch soon enough?

• removed check for ConvTranspose3D on MPS pytorch/pytorch#145366

[valosekj]

Exactly! IIUC, once pytorch/pytorch#145366 is merged, then we'll be able to use the master pytorch branch directly.