Skip to content

Add experimental MPS support #74

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 7 commits into
base: main
Choose a base branch
from
Draft

Add experimental MPS support #74

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
7a9dc13
Add prototype of MPS support
C-Achard May 3, 2024
6f1a646
Fix checking for MPS
C-Achard May 3, 2024
f89ca8e
Update model_framework.py
C-Achard May 3, 2024
3a0dd11
Add fallback environment variable
C-Achard May 3, 2024
822826f
Set fallback mode for MPS directly in workers
C-Achard May 3, 2024
a1077bc
Merge branch 'main' into cy/mps-support
C-Achard May 3, 2024
9d81478
Merge branch 'main' into cy/mps-support
C-Achard May 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions napari_cellseg3d/code_models/model_framework.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
from typing import TYPE_CHECKING

import torch
import torch.backends
import torch.backends.mps

if TYPE_CHECKING:
import napari
Expand Down Expand Up @@ -94,6 +96,17 @@ def __init__(
available_devices = ["CPU"] + [
f"GPU {i}" for i in range(torch.cuda.device_count())
]
from napari_cellseg3d.utils import _is_mps_available

try:
if _is_mps_available(torch):
from os import environ

environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
Copy link
Collaborator Author

@C-Achard C-Achard May 3, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It ought to be set globally as soon as MPS is available... Maybe if I move it; anyway this is not ideal, I don't want you to lose too much time in tests. Does server 6 have MPS ?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it does, but let's focus on the other items first (unifying naming to wnet3d, etc :).

available_devices.append("MPS (beta)")
except Exception as e:
logger.error(f"Error while checking MPS availability : {e}")

self.device_choice = ui.DropdownMenu(
available_devices,
parent=self,
Expand Down Expand Up @@ -345,6 +358,8 @@ def check_device_choice(self):
elif "GPU" in choice:
i = int(choice.split(" ")[1])
device = f"cuda:{i}"
elif choice == "MPS (beta)": # TODO : check if MPS is available
device = "mps"
else:
device = self.get_device()
logger.debug(f"DEVICE choice : {device}")
Expand Down
16 changes: 16 additions & 0 deletions napari_cellseg3d/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import napari
import numpy as np
import pkg_resources
import torch
from monai.transforms import Zoom
from numpy.random import PCG64, Generator
Expand Down Expand Up @@ -645,3 +646,18 @@ def fraction_above_threshold(volume: np.array, threshold=0.5) -> float:
f"non zero in above_thresh : {np.count_nonzero(above_thresh)}"
)
return np.count_nonzero(above_thresh) / np.size(flattened)


def _is_mps_available(torch):
available = False
if pkg_resources.parse_version(
torch.__version__
) >= pkg_resources.parse_version("1.12"):
LOGGER.debug("Torch version is 1.12 or higher, compatible with MPS")
if torch.backends.mps.is_available():
LOGGER.debug("MPS is available")
if torch.backends.mps.is_built():
LOGGER.debug("MPS is built")
available = True

return available
Loading