GPU support for Mac M1?

[alix-pham]

Hi :)

I know the README says GPU speed can only be used with NVIDIA GPUs (Linux and Windows), I was wondering if you were planning on providing support for Mac M1 GPUs in the future?

I found this article, would it work? (I guess not because you use CUDA and CUDA cannot work with M1?)

Thank you very much in advance!

Best,
Alix

[kevinjohncutler]

@alix-pham I will definitely get Apple Silicon GPU support working. It should be as simple as changing the device type from cuda in several places, probably not more than a few hours of work once I get the chance. In my testing so far, the much more difficult thing is getting a conda environment installed that has all the dependencies built for arm64... so releasing an environment is also on my to-do list.

[alix-pham]

Thank you very much @kevinjohncutler! Looking forward to it.

[kevinjohncutler]

@alix-pham I got it working on my M2 Macbook Air. About 2-4x speed improvement on the built-in test images in the GUI. Note that this is for evaluation only, I have not tested training yet. Most of the changes are in the cellpose backend, but pulling the omnipose changes and running pip install -e . should fetch the up to date cellpose changes as well. I need to post an environment file for apple silicon installs, but if you are able to help test, that would be great.

[alix-pham]

Thank you so much @kevinjohncutler! I would say I am not familiar enough with those things to help you test it, unfortunately... I am not even sure I understand what I should do to try GPU support on my Mac with the information you provided... 😬
I updated omnipose using pip install git+https://github.com/kevinjohncutler/omnipose.git but when running the segmentation I get that output (the same as usual)

2022-10-20 11:04:18,530 [INFO] TORCH GPU version not installed/working.
>>> GPU activated? 0
2022-10-20 11:04:18,531 [INFO] >>bact_phase_omni<< model set to be used
2022-10-20 11:04:18,532 [INFO] >>>> using CPU

Is there something else I should do?
Where should I run pip install -e .? I did not clone the repo on my computer (and as I am not sure what I'm supposed to do with it, I didn't run it yet).

Thanks in advance!

PS: Congrats on your Nature Methods paper 🥳

[kevinjohncutler]

Thanks @alix-pham! I see, I figured since you were requesting mac GPU support you knew what a world of pain you were getting yourself into haha. I may finally have time this weekend to put together a conda environment and installation instructions to make it relatively painless. It's possible that the macOS GUI executable will also 'just work', but I need to compile a new version. The issue you are running into is just to do with the dependencies, but the conda environment will aim to solve that.

[alix-pham]

No, sorry! I only want the process to be faster, as we are processing big movies, and using the GPU should make it faster.
Though I am not using the GUI because I'm complementing the segmentation with a tracking pipeline; I figured it would be easier that way.
Thank you very much!

[mccruz07]

Hi @kevinjohncutler,

Thanks for all your work! I can confirm the GPU support for Mac M1/2 with your modifications, but unfortunately I'm receiving an Attribute error while trying to train a new model.

I tested on M1 mac:
Python 3.10.4

and Windows 11
Python = 3.8.4
pytorch = 1.11.0
cudatoolkit = 11.3.1

Here is the report:

> python -m omnipose --train --use_gpu --dir ./Documents/omni --mask_filter _masks --n_epochs 100 --pretrained_model None --learning_rate 0.1 --diameter 0 --batch_size 16 --RAdam
!NEW LOGGING SETUP! To see cellpose progress, set --verbose
No --verbose => no progress or info printed
2022-11-16 16:02:24,512 [INFO] ** TORCH GPU version installed and working. **
2022-11-16 16:02:24,512 [INFO] >>>> using GPU
Omnipose enabled. See Omnipose repo for licencing details.
2022-11-16 16:02:24,512 [INFO] Training omni model. Setting nclasses=4, RAdam=True
2022-11-16 16:02:24,514 [INFO] not all flows are present, will run flow generation for all images
2022-11-16 16:02:24,515 [INFO] training from scratch
2022-11-16 16:02:24,515 [INFO] median diameter set to 0 => no rescaling during training
2022-11-16 16:02:24,601 [INFO] No precomuting flows with Omnipose. Computed during training.
2022-11-16 16:02:24,608 [INFO] >>> Using RAdam optimizer
2022-11-16 16:02:24,608 [INFO] >>>> training network with 2 channel input <<<<
2022-11-16 16:02:24,608 [INFO] >>>> LR: 0.10000, batch_size: 16, weight_decay: 0.00001
2022-11-16 16:02:24,608 [INFO] >>>> ntrain = 2
2022-11-16 16:02:24,608 [INFO] >>>> nimg_per_epoch = 2
/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/core.py:1105: UserWarning: The operator 'aten::linalg_vector_norm' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/_temp/anaconda/conda-bld/pytorch_1668586478573/work/aten/src/ATen/mps/MPSFallback.mm:11.)
  denom = torch.multiply(torch.linalg.norm(x,dim=1),torch.linalg.norm(y,dim=1))+eps
2022-11-16 16:02:33,471 [INFO] Epoch 0, Time  8.9s, Loss 4.7680, LR 0.1000
2022-11-16 16:02:34,080 [INFO] saving network parameters to /Users/mcruz/Documents/omni/models/cellpose_residual_on_style_on_concatenation_off_omni_nclasses_4_omni_2022_11_16_16_02_24.602007
Traceback (most recent call last):
  File "/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/site-packages/omnipose/__main__.py", line 3, in <module>
    main(omni_CLI=True)
  File "/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/__main__.py", line 476, in main
    cpmodel_path = model.train(images, labels, train_files=image_names,
  File "/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/models.py", line 1045, in train
    model_path = self._train_net(train_data, train_labels,
  File "/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/core.py", line 1057, in _train_net
    self.net.save_model(file_name)
  File "/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'DataParallel' object has no attribute 'save_model'  

Best,
Mario

[kevinjohncutler]

@mccruz07 Thanks for the report! I have not tried training on apple silicon yet, but it looks like that might be a simple fix. I'll look into it in the next week.

[kevinjohncutler]

@mccruz07 Turns out all GPU training was broken due to a recent change I made to fix a bug for CPU training. I fixed it now in cellpose-omni v0.7.3. I will test it on an M2 mac in the next couple days, but let me know if you get a chance to test it earlier.

[mccruz07]

@kevinjohncutler Thank you! But know I'm reciving the follow error:

/Users/mcruz/opt/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/core.py:1111: UserWarning: The operator 'aten::linalg_vector_norm' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications.

[kevinjohncutler]

Update: with torch 1.13.1, training is working on Apple Silicon =D @mccruz07 I am still getting that warning, but no errors. Really rough benchmark on a small (5-image) dataset: Titan RTX takes 34.1s for the first 100 epochs, M2 GPU takes 95.7s, so about 2.8x slower. For reference, CPU training on my Ubuntu machine with a core i9 9900K is 161.3s (1.7x slower than M2 GPU) and M2 CPU takes 277.6s (2.9x slower than M2 GPU).

I'm thinking about getting a Mac Studio to have much more VRAM than any consumer NVIDIA card can offer... @michaels10, what config do you have?

[michaels10]

Hi -- I have Mac M1 Ultra; that said, I just tried GPU and it doesn't seem to be working for me; I get a rather uneventful
2023-01-12 17:59:42,705 [INFO] TORCH GPU version not installed/working. error message.

I used pip install -e . on the cloned repo. I'm running python -m cellpose --train --pretrained_model bact_omni --use_gpu --chan 0 --dir "/Users/michaelsandler/Documents/experiments/omnipose-test/training-data" --n_epochs 100 --learning_rate 0.1 --verbose command to run things. Pytorch version 1.12, which I believe in their versioning scheme is greater than 1.4?

Minimal set to reproduce same as in the other bug.

[kevinjohncutler]

Thanks @michaels10, good to know. 64 or 128GB of RAM? In addition to the cellose_omni bug which I hope is now fixed for you (it should now download v0.8.0 or higher), it's pobably that your conda environment is not set up for pytorch on M1 - you are right, torch 1.13.1 is what you want. I will update this repo with an environment file for Macs.

[kevinjohncutler]

Ok, try out omnipose_mac_environment.yml. I installed it with

conda env create --name omnipose --file /Volumes/DataDrive/omnipose_mac_environment.yml
conda activate omnipose
pip install git+https://github.com/kevinjohncutler/omnipose.git
pip install git+https://github.com/kevinjohncutler/cellpose-omni.git

To my amazement, it worked the first time around. However, I have some notes from my first attempts getting this to work months back, and it is possible that some dependencies actually need to be compiled from source and my conda environment is just using the right versions from the base environment... we shall see once more people try this out.

[michaels10]

Tried it, alas to no avail. **
PyTorch is version 1.13.1 -- also, my computer is the 128GB model.

The only potentially relevant warning I get is:

/Users/michaelsandler/opt/anaconda3/envs/omnipose/lib/python3.9/runpy.py:127: RuntimeWarning: 'cellpose_omni.__main__' found in sys.modules after import of package 'cellpose_omni', but prior to execution of 'cellpose_omni.__main__'; this may result in unpredictable behaviour

**Weirdly, someone in our lab with the exact same setup is having no problems. Maybe something is wrong with my base environment.