BUG: MPS support?

[lara-bonney]

Hi,

I was trying to run Moose for the first time as a trial on my MacBook, however, I get a similar error to that seen in bug reports #125, #91 + #78. Whereby it looks like the prediction happens but no output files are saved in output_dir which then causes the process to error out. I'm currently just running this over one participant directory as a test but hoping to use throughout our project. My MacBook has Apple Silicon GPU and this appears to be detected by Moose at the start. The model is already downloaded in the run below as I've tried before.

This very well could be within the realm of user error and I apologise if so! Although it seems very similar to a few other bugs that were reported and I didn't quite understand from those threads how those were resolved. Thank you very much in advance for any help you can provide.

Things tried and achieved same error:

• Other models from the set.
• Creating a fresh environment and restarting the whole process.
• Restarting my Mac to clear some RAM (there is plenty available)

I ran the code over a directory called MOOSE-TEST which contained one participant S3 with a DICOM series in a folder called CT-BODY.

Terminal input/output:

(moose-env2) (base) lara@LarasMaook92586 ~ % moosez -d /Users/lara/Downloads/MOOSE-TEST -m clin_ct_organs

Apple MPS backend is available. Predictions will be run on Apple Silicon GPU.

🌐 MODEL DOWNLOAD:

A local instance of Dataset123_Organs has been detected.

🔍 STANDARDIZING INPUT DATA TO NIFTI:

Processing S3... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00 Standardization complete. Number of moose compliant subjects: 1 out of 1

🔮 PREDICT:

⠙ [1/1] Running prediction for S3 using clin_ct_organs...Traceback (most recent call last): File "/Users/lara/moose-env2/bin/moosez", line 8, in sys.exit(main()) File "/Users/lara/moose-env2/lib/python3.10/site-packages/moosez/moosez.py", line 198, in main predict.predict(model_name, input_dir, output_dir, accelerator) File "/Users/lara/moose-env2/lib/python3.10/site-packages/moosez/predict.py", line 79, in predict postprocess(original_image_files[0], output_dir, model_name) File "/Users/lara/moose-env2/lib/python3.10/site-packages/moosez/predict.py", line 132, in postprocess predicted_image = file_utilities.get_files(output_dir, '.nii.gz')[0] IndexError: list index out of range

Contents of output_dir Nifty file of CT is saved, when visualised this appears fine, no artefacts etc.

Segmentations and stats folders are created. The segmentation folder contains three files dataset.json, plans.json and predict_from_raw_data_args.json, while the stats folder is empty.

[LalithShiyam]

Hi @lara-bonney, I am away from my m2 Mac! Will need 3 more hours to fix this :)! Sorry about that. I will look into it.

[lara-bonney]

There's no rush! Thank you very much for the quick reply!

[mprires]

Hi @lara-bonney. I was able to reproduce your error. Unfortunately nnUNet, the segmentation used inside MOOSE doesn't work for 3d segmentation with MPS. I'm fixing this and I'll let you know when it's online

[LalithShiyam]

@mprires I think this should be supported by PyTorch now? I remember there was a patch that was pushed in - which means PyTorch can use MPS, which means nnunetv2 can, which means we can. Kindly confirm if this is not done. If that's the case, I will merge your PR.

@lara-bonney thank you for your patience.

Cheers, Lalith

[mprires]

@LalithShiyam According to nnUNet documentation it is not supproted for 3D segmentation

[LalithShiyam]

@lara-bonney It is not straightforward to use MPS as of now. And its a bit tricky to get it work. I will remove this as of now, so that mac users can atleast run it on CPU. We will soon come up with a fix. You can upgrade moosez 'pip install --upgrade moosez' so that you can run things on your mac.

Or just to be safe - please create a new environment and make a fresh install of moosez. 2.4.11 is your friend.

[lara-bonney]

@LalithShiyam and @mprires

Thank you so much for the very prompt response. I have tried the updated version and it is hasn't yet finished but it appears to be working, I will leave it running overnight tonight - thank you again!

[LalithShiyam]

No worries - it will definitely take a while, since its CPU. Things are faster with GPU. Sorry about that.

[lara-bonney]

No need to apologise at all! I was expecting it to take a while on CPU, I just wanted to let you know it looks like it's working :)

[LalithShiyam]

May I close this issue then :)

[lara-bonney]

Yes, it worked overnight! Thank you again for your help.

[LalithShiyam]

Hi @lara-bonney, I had to re-open your issue, because it kept me awake and I would like to also use my M2 chips. So I managed to get it working, are you still interested in this? If yes, I can create a discussion board on how to get this running in an MPS Mac. And frankly, the speed up is impressive. I got to segment the organs in 2.3 min per dataset - which is way better than CPU and not so bad compared to a GPU. Let me know if you are interested and I am happy to create a discussion board on how to install.

[lara-bonney]

Hi @LalithShiyam I am definitely still interested in this! That's a very impressive improvement in speed. Thank you for looking into this further

[LalithShiyam]

@lara-bonney fantastic - can you please tell me your Macbook's hardware spec? So that we can benchmark?

[LalithShiyam]

@lara-bonney These are the steps, please follow them without any deviations and it should work.

1. Create the python3.10 env: python3.10 -m venv moosez-mps
2. And go inside the environment: source moosez-mps/bin/activate
3. Install moosez from github main branch: pip install moosez
4. Install this special torch branch which contains the convTranspose3D working for MPS: pip install git+https://github.com/pytorch/pytorch.git@3c61c525694eca0f895bb01fc67c16793226051a
5. Now run moose as you usually do and hopefully it will work! Keep me posted and also let me know how long it took - I am very keen on benchmarking this.
6. Also ignore the leak issues at the end, this should be fine with a stable release of torch - I think!

GPU utilisation

Time taken for 3 datasets

[lara-bonney]

@LalithShiyam I'll have a go at this tomorrow, thank you for the instructions.

[LalithShiyam]

Thank you! 🥳 Hope it works!

[lara-bonney]

Hi @LalithShiyam, It worked - thank you!! It's marginally slower than you saw but my Mac is M2 Max as opposed to an M1 ultra so I think that would be expected. Also no leak issues as you suspected. Hopefully all the info you need is below, I did this on 4 CT scans I have for my project, is there a test dataset you would like me to run for your benchmarking?

Mac Specs: Model Name: MacBook Pro Chip: Apple M2 Max Total Number of Cores: 12 (8 performance and 4 efficiency) Memory: 32 GB

GPU Utilisation

Time taken for 4 datasets:

[LalithShiyam]

@lara-bonney Fantastic - this is great news, what is the matrix size of your CT images? Also I will reach out to you with a 'benchmark' dataset. We don't have one yet, but I can figure out something quick tomorrow. Many thanks for all your effort. This is a nice bump in performance!

[lara-bonney]

@LalithShiyam Matrix size is 512x512, axial length varies. It's impressively fast, thank you for getting this to work!

[LalithShiyam]

I am happy it's working well. Finally can this issue can be closed for real? 🥇

[lara-bonney]

Yes of course!