def __init__(self, id_or_device):
_type, id = None, None
if isinstance(id_or_device, int):
id = id_or_device
elif id_or_device == "cpu":
_type, id = id_or_device, -1
elif isinstance(id_or_device, Device):
_type, id = id_or_device.type, id_or_device.id
elif env.cupy_available() and isinstance(id_or_device, cp.cuda.Device):
_type, id = "cuda", id_or_device.id
elif env.sigpy_available() and isinstance(id_or_device, sp.Device):
id = id_or_device.id
elif env.torch_available() and isinstance(id_or_device, torch.device):
_type, id = id_or_device.type, id_or_device.index
if id is None:
if _type == "cuda":
id = torch.cuda.current_device()
elif _type == "cpu":
id = -1
else:
raise ValueError(f"Unsupported device type: {_type}")
else:
raise ValueError(
f"Accepts int, Device, cupy.cuda.Device, or torch.device"
f"got {id_or_device}")
assert id >= -1
if _type is None:
_type = "cpu" if id == -1 else "cuda"
cpdevice = None
if id != -1:
if env.cupy_available():
cpdevice = cp.cuda.Device(id)
else:
raise ValueError(
"cupy not installed, but set device {}.".format(id))
self._type = _type
self._id = id
self._cpdevice = cpdevice
def _process_mask(self, mask, y: MedicalVolume):
"""Process mask into appropriate shape."""
arr_types = (np.ndarray,
cp.ndarray) if env.cupy_available() else (np.ndarray, )
if isinstance(mask, arr_types):
mask = y._partial_clone(volume=mask, headers=None)
elif not isinstance(mask, MedicalVolume):
raise TypeError("`mask` must be a MedicalVolume or ndarray")
mask = mask.reformat_as(y)
if not mask.is_same_dimensions(y, defaults.AFFINE_DECIMAL_PRECISION):
raise RuntimeError("`mask` and `y` dimension mismatch")
return mask > 0
def __eq__(self, other):
if other == -1:
# other integers are not compared as self.type may be subject to change
return self.id == other
elif isinstance(other, Device):
return self.type == other.type and self.id == other.id
elif env.cupy_available() and isinstance(other, cp.cuda.Device):
return self.type == "cuda" and self.id == other.id
elif env.sigpy_available() and isinstance(other, sp.Device):
try:
return self.spdevice == other
except RuntimeError:
return False
elif env.torch_available() and isinstance(other, torch.device):
return self.ptdevice == other
else:
return False
def get_array_module(array):
"""Gets an appropriate module from :mod:`numpy` or :mod:`cupy`.
This is almost equivalent to :func:`cupy.get_array_module`. The differences
are that this function can be used even if cupy is not available.
Adapted from :mod:`sigpy`.
Args:
array: Input array.
Returns:
module: :mod:`cupy` or :mod:`numpy` is returned based on input.
"""
if env.cupy_available():
return cp.get_array_module(array)
else:
return np
def _extract_input_array_ufunc(self, input, device=None):
if device is None:
device = self.device
device_err = "Expected device {} but got device ".format(device) + "{}"
if isinstance(input, Number):
return input
elif isinstance(input, np.ndarray):
if device != cpu_device:
raise RuntimeError(device_err.format(cpu_device))
return input
elif env.cupy_available() and isinstance(input, cp.ndarray):
if device != input.device:
raise RuntimeError(device_err.format(Device(input.device)))
return input
elif isinstance(input, MedicalVolume):
if device != input.device:
raise RuntimeError(device_err.format(Device(input.device)))
assert self.is_same_dimensions(input, err=True)
return input._volume
else:
return NotImplemented
"""Functions and classes for getting and setting computing devices.
"""
import numpy as np
from dosma.utils import env
if env.cupy_available():
import cupy as cp
if env.sigpy_available():
import sigpy as sp
if env.torch_available():
import torch
__all__ = ["Device", "get_device", "to_device"]
class Device(object):
"""Device class.
This class extends ``cupy.Device`` and can also be used to interface with
``torch.Device`` and ``sigpy.Device``. This class contains a device type
('cpu' or 'cuda') and optional device ordinal (i.e. the id) for the device type.
This class can also be constructed using only the ordinal id, where id >= 0
representing the id-th GPU, and id = -1 representing CPU. cupy must be installed to use GPUs.
The array module for the corresponding device can be obtained via ``device.xp``.
Similar to cupy.Device, the Device object can be used as a context:
>>> device = Device(1) # gpu 1
>>> xp = device.xp # xp is cupy.