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))