Adding GPU support to the GUI in Apple Silicon

[PikaPei]

Hello Cellpose team!

I noticed that M1 GPU support is in the command-line mode in #523, but not yet in GUI.

This PR checks if the platform is 'arm', the torch.device will use 'mps'.
I have tested it on my M2 macbook, and hope it's useful.

[sophiamaedler]

I tried installing cellpose from @PikaPei fork to see if I can get a GUI running with GPU support in Apple Silicon but I encountered the following error:

NET ERROR: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

Executing the same segmentation in a python script I got the following traceback:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/models.py", line 216, in eval
    diams, _ = self.sz.eval(x, channels=channels, channel_axis=channel_axis, invert=invert, batch_size=batch_size, 
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/models.py", line 889, in eval
    diam, diam_style = self.eval(x[i], 
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/models.py", line 925, in eval
    masks = self.cp.eval(x, 
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/models.py", line 552, in eval
    masks, styles, dP, cellprob, p = self._run_cp(x, 
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/models.py", line 652, in _run_cp
    outputs = dynamics.compute_masks(dP[:,i], cellprob[i], niter=niter, cellprob_threshold=cellprob_threshold,
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/dynamics.py", line 746, in compute_masks
    mask = remove_bad_flow_masks(mask, dP, threshold=flow_threshold, use_gpu=use_gpu, device=device)
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/dynamics.py", line 601, in remove_bad_flow_masks
    merrors, _ = metrics.flow_error(masks, flows, use_gpu, device)
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/metrics.py", line 282, in flow_error
    dP_masks = dynamics.masks_to_flows(maski, use_gpu=use_gpu, device=device)
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/dynamics.py", line 277, in masks_to_flows
    mu, mu_c = masks_to_flows_device(masks, device=device)
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/dynamics.py", line 152, in masks_to_flows_gpu
    mu = _extend_centers_gpu(neighbors, centers, isneighbor, Ly, Lx, 
  File "/Users/sophia/mambaforge/envs/cellpose2/lib/python3.9/site-packages/cellpose/dynamics.py", line 71, in _extend_centers_gpu
    T = torch.zeros((nimg,Ly,Lx), dtype=torch.double, device=device)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

By changing dtype=torch.double to dtype=torch.float32 in dynamics.py it seems to fix the issue I can successfully segment an image both in the GUI and in a python terminal using my GPU.

I am unfortunately not sure how to add this as an additional commit to this PR.

[OratHelm]

As I explained there, I had to make a few more changes to get it to work properly for me. I've proposed a fork, but I think it will be easy enough to integrate into the main branch.
My main problem is that I can't get the training to work. It runs on the GPU, without errors, but the model generated is of no use. See here for more details. If anyone has any ideas, that would be great.

[carsen-stringer]

thanks everyone for your help with this, we've merged in @OratHelm 's fork and it should be working (training included), let us know if not