def _guess_device_from_array_module(xp):
"""Returns a plausible device from array module
.. warning::
There can be multiple devices for a module
"""
if xp is cuda.cupy:
return cuda.GpuDevice(cuda.Device())
elif xp is chainerx:
return _chainerx.ChainerxDevice(chainerx.get_default_device())
else:
# Cannot detect intel64, because xp of intel64 is numpy.
return _cpu.CpuDevice()
def get_device(device_spec: types.DeviceSpec) -> Device:
"""Returns a device object.
Args:
device_spec (object): Device specifier.
If a :class:`chainer.backend.Device` instance is given, it is
returned intact. Otherwise the following values are supported:
* ChainerX devices
* A string representing a device.
(ex. ``'native:0'``, ``'native'``)
* A :class:`chainerx.Device` object.
* CuPy
* A string starts with ``'@cupy:'``.
(ex. ``'@cupy:0'``)
* A :class:`cupy.cuda.Device` object.
* NumPy
* The string ``'@numpy'``.
* NumPy with Intel Architecture
* The string ``'@intel64'``.
"""
if isinstance(device_spec, Device):
return device_spec
if isinstance(device_spec, cuda._integer_types):
return _get_device_cupy_or_numpy(device_spec)
if chainerx.is_available() and isinstance(device_spec, chainerx.Device):
return _chainerx.ChainerxDevice(device_spec)
if cuda.available and isinstance(device_spec, cuda.Device):
return cuda.GpuDevice(device_spec)
if isinstance(device_spec, six.string_types):
# '-1', '0', '1', ...
try:
int_device_spec = int(device_spec)
except ValueError:
pass
else:
return _get_device_cupy_or_numpy(int_device_spec)
if device_spec.startswith('@'):
# '@module:...'
mod_name, colon, precise_spec = device_spec[1:].partition(':')
if mod_name == 'numpy':
if not colon:
return _cpu.CpuDevice()
elif mod_name == 'cupy':
if colon:
return cuda.GpuDevice.from_device_id(int(precise_spec))
elif mod_name == 'intel64':
if not colon:
return intel64.Intel64Device()
raise ValueError(
'Device specifiers starting with \'@\' must be followed by'
' a module name and depending on the module, module specific'
' precise device specifiers. Actual: {}'.format(device_spec))
else:
# String device specifier without '@' prefix is assumed to be a
# ChainerX device.
if not chainerx.is_available():
raise RuntimeError(
'Tried to parse ChainerX device specifier \'{}\', '
'but ChainerX is not available. '
'Note that device specifiers without \'@\' prefix are '
'assumed to be ChainerX device '
'specifiers.'.format(device_spec))
return _chainerx.ChainerxDevice(chainerx.get_device(device_spec))
raise TypeError(
'Device specifier must be a backend.Device, cuda.Device,'
' chainerx.Device, integer or a string. Actual: {}'.format(
type(device_spec)))
def get_device(device_spec):
# type: (types.DeviceSpec) -> Device
"""Returns a device object.
Args:
device_spec (object): Device specifier.
If a :class:`chainer.backend.Device` instance is given, it is
returned intact. Otherwise the following values are supported:
* ChainerX devices
* A string representing a device.
(ex. ``'native:0'``, ``'native'``)
* A :class:`chainerx.Device` object.
* CuPy
* A string starts with ``'@cupy:'``.
(ex. ``'@cupy:0'``)
* A :class:`cupy.cuda.Device` object.
* NumPy
* The string ``'@numpy'``.
* NumPy with Intel Architecture
* The string ``'@intel64'``.
"""
if isinstance(device_spec, Device):
return device_spec
if isinstance(device_spec, cuda._integer_types):
return _get_device_cupy_or_numpy(device_spec)
if chainerx.is_available() and isinstance(device_spec, chainerx.Device):
return _chainerx.ChainerxDevice(device_spec)
if cuda.available and isinstance(device_spec, cuda.Device):
return cuda.GpuDevice(device_spec)
if isinstance(device_spec, six.string_types):
# '-1', '0', '1', ...
try:
int_device_spec = int(device_spec)
except ValueError:
pass
else:
return _get_device_cupy_or_numpy(int_device_spec)
if device_spec.startswith('@'):
# '@module:...'
mod_name, colon, precise_spec = device_spec[1:].partition(':')
if mod_name == 'numpy':
if not colon:
return _cpu.CpuDevice()
elif mod_name == 'cupy':
if colon:
return cuda.GpuDevice.from_device_id(int(precise_spec))
elif mod_name == 'intel64':
if not colon:
return intel64.Intel64Device()
elif chainerx.is_available():
return _chainerx.ChainerxDevice(chainerx.get_device(device_spec))
raise ValueError('Invalid device specifier: {}'.format(device_spec))