def test_sigpy_cupy(self):
import sigpy as sp
assert Device(0) == sp.Device(0)
device = Device(0)
assert device.spdevice == sp.Device(0)
def test_device_gpu(self):
import cupy as cp
mv = MedicalVolume(np.ones((10, 20, 30)), self._AFFINE)
mv_gpu = mv.to(Device(0))
assert mv_gpu.device == Device(0)
assert isinstance(mv_gpu.volume, cp.ndarray)
assert isinstance(mv_gpu.affine, np.ndarray)
assert mv_gpu.is_same_dimensions(mv)
assert cp.all((mv_gpu + 1).volume == 2)
assert cp.all((mv_gpu - 1).volume == 0)
assert cp.all((mv_gpu * 2).volume == 2)
assert cp.all((mv_gpu / 2).volume == 0.5)
assert cp.all((mv_gpu > 0).volume)
assert cp.all((mv_gpu >= 0).volume)
assert cp.all((mv_gpu < 2).volume)
assert cp.all((mv_gpu <= 2).volume)
ornt = tuple(x[::-1] for x in mv_gpu.orientation[::-1])
mv2 = mv_gpu.reformat(ornt)
assert mv2.orientation == ornt
mv_cpu = mv_gpu.cpu()
assert mv_cpu.device == Device(-1)
assert mv_cpu.is_identical(mv)
with self.assertRaises(RuntimeError):
mv_gpu.save_volume(
os.path.join(self._TEMP_PATH, "test_device.nii.gz"))
def test_sigpy(self):
import sigpy as sp
assert Device(-1) == sp.cpu_device
assert Device(sp.cpu_device) == sp.cpu_device
device = Device(-1)
assert device == sp.cpu_device
assert device.spdevice == sp.cpu_device
def test_torch(self):
import torch
pt_device = torch.device("cpu")
assert Device(pt_device) == cpu_device
dm_device = Device(-1)
assert dm_device == pt_device
assert dm_device.ptdevice == pt_device
def test_cupy(self):
import cupy as cp
device = Device(0)
assert device.cpdevice == cp.cuda.Device(0)
assert device.type == "cuda"
assert device.index == 0
assert device.xp == cp
assert int(device) == 0
device = Device(cp.cuda.Device(0))
assert device.cpdevice == cp.cuda.Device(0)
assert device.type == "cuda"
assert device.index == 0
def test_basic(self):
assert Device(-1) == cpu_device
assert Device("cpu") == cpu_device
assert cpu_device.xp == np
device = Device(-1)
assert int(device) == -1
assert device.index == -1
assert device.id == -1
assert device == -1
assert device.cpdevice is None
device2 = Device(-1)
assert device2 == device
def test_array_gpu(self):
import cupy as cp
mv = MedicalVolume(np.ones((10, 20, 30)), self._AFFINE)
mv_gpu = mv.to(Device(0))
data = cp.asarray(mv_gpu)
assert cp.shares_memory(data, mv_gpu.volume)
def test_hdf5(self):
shape = (10, 20, 30)
volume = np.reshape(list(range(np.product(shape))), shape)
hdf5_file = os.path.join(self._TEMP_PATH, "unittest.h5")
with h5py.File(hdf5_file, "w") as f:
f.create_dataset("volume", data=volume)
f = h5py.File(hdf5_file, "r")
mv = MedicalVolume(f["volume"], np.eye(4))
assert mv.device == Device("cpu")
assert mv.dtype == f["volume"].dtype
mv2 = mv[:, :, :1]
assert np.all(mv2.volume == volume[:, :, :1])
assert mv2.device == Device("cpu")
assert mv2.dtype == volume.dtype
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
def to(self, device):
"""Move to device.
If on same device, no-op and returns ``self``.
Args:
device: The device to move to.
Returns:
MedicalVolume
"""
device = Device(device)
if self.device == device:
return self
return self._partial_clone(volume=to_device(self._volume, device))
def test_device(self):
mv = MedicalVolume(np.ones((10, 20, 30)), self._AFFINE)
assert mv.to(Device(-1)) == mv
assert mv.cpu() == mv