Improve device support and add support for Apple Silicon chipset (mps)

cebra/helper.py

@@ -18,6 +18,7 @@
 import zipfile
 from typing import List
 
+import pkg_resources
 import requests
 
 import cebra.data
@@ -88,3 +89,14 @@ def download_file_from_zip_url(url, file="montblanc_tracks.h5"):
  except zipfile.error:
  pass
  return pathlib.Path(foldername) / "data" / file
+
+
+def _is_mps_availabe(torch):
+ available = False
+ if pkg_resources.parse_version(
+ torch.__version__) >= pkg_resources.parse_version("1.12"):
+ if torch.backends.mps.is_available():
+ if torch.backends.mps.is_built():
+ available = True
+
+ return available

cebra/integrations/sklearn/cebra.py

@@ -1256,3 +1256,52 @@ def load(cls,
  raise RuntimeError("Model loaded from file is not compatible with "
  "the current CEBRA version.")
  return model
+
+ def to(self, device: Union[str, torch.device]):
+ """Moves the cebra model to the specified device.
+
+ Args:
+ device: The device to move the cebra model to. This can be a string representing
+ the device ('cpu','cuda', cuda:device_id, or 'mps') or a torch.device object.
+
+ Returns:
+ The cebra model instance.
+
+ Example:
+
+ >>> import cebra
+ >>> import numpy as np
+ >>> dataset = np.random.uniform(0, 1, (1000, 30))
+ >>> cebra_model = cebra.CEBRA(max_iterations=10, device = "cuda_if_available")
+ >>> cebra_model.fit(dataset)
+ CEBRA(max_iterations=10)
+ >>> cebra_model = cebra_model.to("cpu")
+ """
+
+ if not isinstance(device, (str, torch.device)):
+ raise TypeError(
+ "The 'device' parameter must be a string or torch.device object."
+ )
+
+ if (not device == 'cpu') and (not device.startswith('cuda')) and (
+ not device == 'mps'):
+ raise ValueError(
+ "The 'device' parameter must be a valid device string or device object."
+ )
+
+ if isinstance(device, str):
+ device = torch.device(device)
+
+ if (not device.type == 'cpu') and (
+ not device.type.startswith('cuda')) and (not device == 'mps'):
+ raise ValueError(
+ "The 'device' parameter must be a valid device string or device object."
+ )
+
+ if hasattr(self, "device_"):
+ self.device_ = device
+
+ self.device = device
+ self.solver_.model.to(device)
+
+ return self

cebra/integrations/sklearn/utils.py

@@ -15,6 +15,8 @@
 import sklearn.utils.validation as sklearn_utils_validation
 import torch
 
+import cebra.helper
+
 
 def update_old_param(old: dict, new: dict, kwargs: dict, default) -> tuple:
  """Handle deprecated arguments of a function until they are replaced.
@@ -114,16 +116,50 @@ def check_device(device: str) -> str:
  device: The device to return, if possible.
 
  Returns:
- Either cuda or cpu depending on {device} and availability in the environment.
+ Either cuda, cuda:device_id, mps, or cpu depending on {device} and availability in the environment.
  """
+
  if device == "cuda_if_available":
  if torch.cuda.is_available():
  return "cuda"
+ elif cebra.helper._is_mps_availabe(torch):
+ return "mps"
  else:
  return "cpu"
- elif device in ["cuda", "cpu"]:
+ elif device.startswith("cuda:") and len(device) > 5:
+ cuda_device_id = device[5:]
+ if cuda_device_id.isdigit():
+ device_count = torch.cuda.device_count()
+ device_id = int(cuda_device_id)
+ if device_id < device_count:
+ return f"cuda:{device_id}"
+ else:
+ raise ValueError(
+ f"CUDA device {device_id} is not available. Available device IDs are 0 to {device_count - 1}."
+ )
+ else:
+ raise ValueError(
+ f"Invalid CUDA device ID format. Please use 'cuda:device_id' where '{cuda_device_id}' is an integer."
+ )
+ elif device == "cuda" and torch.cuda.is_available():
+ return "cuda:0"
+ elif device == "cpu":
  return device
- raise ValueError(f"Device needs to be cuda or cpu, but got {device}.")
+ elif device == "mps":
+ if not torch.backends.mps.is_available():
+ if not torch.backends.mps.is_built():
+ raise ValueError(
+ "MPS not available because the current PyTorch install was not "
+ "built with MPS enabled.")
+ else:
+ raise ValueError(
+ "MPS not available because the current MacOS version is not 12.3+ "
+ "and/or you do not have an MPS-enabled device on this machine."
+ )
+
+ return device
+
+ raise ValueError(f"Device needs to be cuda, cpu or mps, but got {device}.")
 
 
 def check_fitted(model: "cebra.models.Model") -> bool:

cebra/io.py

@@ -75,7 +75,7 @@ def _set_device(self, device):
  return
  if not isinstance(device, str):
  device = device.type
- if device not in ("cpu", "cuda"):
+ if device not in ("cpu", "cuda", "mps"):
  if device.startswith("cuda"):
  _, id_ = device.split(":")
  if int(id_) >= torch.cuda.device_count():

tests/test_sklearn.py

@@ -20,6 +20,7 @@
 import torch
 
 import cebra.data as cebra_data
+import cebra.helper
 import cebra.integrations.sklearn.cebra as cebra_sklearn_cebra
 import cebra.integrations.sklearn.dataset as cebra_sklearn_dataset
 import cebra.models
@@ -823,6 +824,59 @@ def test_save_and_load(action):
  _assert_equal(original_model, loaded_model)
 
 
+def get_ordered_cuda_devices():
+ available_devices = ['cuda']
+ for i in range(torch.cuda.device_count()):
+ available_devices.append(f'cuda:{i}')
+ return available_devices
+
+
+ordered_cuda_devices = get_ordered_cuda_devices() if torch.cuda.is_available(
+) else []
+mps_device = ["mps"] if cebra.helper._is_mps_availabe(torch) else []
+
+
+@pytest.mark.parametrize("device", ['cpu'] + ordered_cuda_devices + mps_device)
+def test_to_device(device):
+
+ # example dataset
+ X = np.random.uniform(0, 1, (1000, 50))
+
+ # Check if CUDA is available
+ if not torch.cuda.is_available():
+ pytest.skip("CUDA not available")
+
+ # Check if device ID exists
+ if device.startswith('cuda') and device not in ordered_cuda_devices:
+ pytest.skip(f"Device ID {device} not available")
+
+ # Create CEBRA model
+ cebra_model = cebra_sklearn_cebra.CEBRA(model_architecture="offset1-model",
+ max_iterations=5,
+ device=device)
+
+ # Train model (time contrastive)
+ cebra_model.fit(X)
+
+ # Move the model to a different device
+ new_device = 'cpu' if device.startswith('cuda') else 'cuda:0'
+ cebra_model.to(new_device)
+
+ # Check that the device has changed
+ assert cebra_model.device == torch.device(new_device)
+
+ # Check that the model parameters are on the correct device
+ assert next(cebra_model.solver_.model.parameters()).device == torch.device(
+ new_device)
+
+ # Check that if you save a model on one device, it stays on same device
+ with tempfile.NamedTemporaryFile(mode="w+b", delete=True) as savefile:
+ cebra_model.save(savefile.name)
+ loaded_model = cebra_sklearn_cebra.CEBRA.load(savefile.name)
+
+ assert cebra_model.device == loaded_model.device
+
+
 @pytest.mark.parametrize("device", ["cpu"] +
  ["cuda"] if torch.cuda.is_available() else [])
 @pytest.mark.parametrize("action", _iterate_actions())