def get_device_spec():
"""Return the device spec in current nesting."""
dev_info = get_device_info()
if dev_info is not None:
return device_spec.DeviceSpec(
dev_info['device_type'],
dev_info['device_index'],
)
else:
cfg = config.config()
return device_spec.DeviceSpec(
cfg.device_type,
cfg.device_index,
)
def get_device(use_default=True):
"""Return the nesting or default device."""
spec = _GLOBAL_DEVICE_STACK.get_default()
if spec is None:
cfg = config.config()
spec = device_spec.DeviceSpec(cfg.device_type, cfg.device_index)
return spec
def device(device_type, device_index=0):
"""Context-manager to nest the device.
Examples:
```python
with dragon.device('cuda', 0):
x = dragon.constant(1)
```
Parameters
----------
device_type : str
The device type.
device_index : int, optional, default=0
The device index.
Returns
-------
dragon.DeviceSpec
The current nesting device spec.
"""
device_type = device_type.lower()
if device_type not in mapping.DEVICE_STRING_TO_DEVICE_TYPE:
raise ValueError('Unsupported device type: ' + device_type)
device_type = mapping.DEVICE_STRING_TO_DEVICE_TYPE[device_type]
spec = device_spec.DeviceSpec(device_type, device_index)
return _GLOBAL_DEVICE_STACK.get_controller(spec)
def _steal_grad(ws, source):
"""Steal the grad from backend."""
impl = ws.GetTensor(source.id + '_grad')
if impl is None:
return None
device = device_spec.DeviceSpec(*impl.device)
return EagerTensor(impl=impl, device=device)
def __call__(self, *args, **kwargs):
"""Call the compiled executables."""
if self.executables is None:
# Graph is not created on the first call.
# Compile the executables from the python function.
inputs = []
input_signature = self.input_signature
with context.name_scope('${%d}' % id(self)):
for i in range(self._function_spec.num_inputs):
name, shape, dtype = 'Input:%d' % i, None, None
if input_signature is not None:
if i >= len(input_signature):
raise ValueError(
'When <input_signature> is provided, '
'only define arguments covered by it.\n'
'Got %d signature(s) and %d argument(s).' %
(len(input_signature),
self._function_spec.num_inputs))
shape = input_signature[i].shape
dtype = input_signature[i].dtype
inputs.append(Tensor(shape, dtype, name).constant())
with context.name_scope('${%d}' %
id(self)), eager_context.graph_mode():
returns = nest.flatten(self._python_function(*inputs))
outputs, dummies = [], []
for obj in returns:
if isinstance(obj, Tensor):
outputs.append(obj)
else:
dummies.append(obj)
executables = [function_lib.create_function(outputs=outputs)]
for obj in dummies:
if isinstance(obj, optimizer.Optimizer):
executables.append(
function_lib.create_function(optimizer=obj))
self.inputs = inputs
self.outputs = returns
self.executables = executables
# In this case, we have compiled executables.
# Notify the backend to run directly.
executables = self.executables
inputs, kwargs = self.canonicalize_inputs(*args, **kwargs)
current_ws = workspace.get_workspace()
for input, value in zip(self.inputs, inputs):
current_ws.feed_tensor(input, value)
executables[0](return_outputs=False, **kwargs)
[func(return_outputs=False) for func in executables[1:]]
outputs = []
for output in self.outputs:
if isinstance(output, Tensor):
impl = current_ws.GetTensor(output.id)
device = device_spec.DeviceSpec(*impl.device)
outputs.append(EagerTensor(impl=impl, device=device))
else:
outputs.append(output)
return outputs[0] if len(outputs) == 1 else outputs
def device(self):
"""Return the tensor device.
Returns
-------
dragon.DeviceSpec
The device.
"""
return device_spec.DeviceSpec(*self._impl.device)
def __init__(self, model, device, **kwargs):
"""Create a ``BackendRep``.
Parameters
----------
model : str
The path of onnx model file.
device : onnx.Device
The executing device.
"""
if not isinstance(device, Device):
device = Device(device)
graph_str = workspace.get_workspace().PrepareONNXModel(model)
graph_def = dragon_pb2.GraphDef()
graph_def.ParseFromString(graph_str)
if device.type == DeviceType.CPU:
device_type, device_index = 'cpu', 0
elif device.type == DeviceType.CUDA:
device_type, device_index = 'cuda', device.device_id
else:
raise ValueError('Unsupported device type: ' + device.type)
with context.device(device_type, device_index):
self._function = function_lib.Function(name='ONNXGraph') \
.import_from(graph_def)
self._input_dict = collections.OrderedDict([
(impl.name,
EagerTensor(impl=impl,
device=device_spec.DeviceSpec(device_type,
device_index)))
for impl in self._function.inputs
])
self._output_dict = collections.OrderedDict([
(impl.name,
EagerTensor(impl=impl,
device=device_spec.DeviceSpec(device_type,
device_index)))
for impl in self._function.outputs
])
def host_dlpack(self):
"""Return the dlpack tensor wrapping host buffer.
Returns
-------
PyCapsule
The dlpack tensor object.
"""
if self._host_tensor is None:
spec = device_spec.DeviceSpec('cpu')
self._host_opt = spec.to_proto(serialized=True)
default_ws = workspace.get_workspace()
impl = default_ws.create_tensor(scope='DLPack')
impl.FromPointer(self._shape, self._dtype,
self._host_opt, self.host_buffer.ctypes.data)
self._host_tensor = Tensor(impl=impl, deleter=default_ws._handle_pool)
return self._host_tensor._impl.ToDLPack(self._host_opt, True)
def host_dlpack(self):
"""Return the dlpack tensor wrapping host buffer.
Returns
-------
PyCapsule
The dlpack tensor object.
"""
if self._host_tensor is None:
spec = device_spec.DeviceSpec('cpu')
self._host_opt = spec.to_proto(serialized=True)
current_ws = workspace.get_workspace()
tensor = EagerTensor(device=spec) # Hack the constructor.
tensor._gc = current_ws.collectors.TENSOR
tensor._impl = current_ws.create_tensor(
tensor._gc.alloc('${DLPACK}')).FromPointer(
self._shape, self._dtype, self._host_opt,
self.host_buffer.ctypes.data)
self._host_tensor = tensor
return self._host_tensor._impl.ToDLPack(self._host_opt, True)
def from_dlpack(dlpack):
"""Create a tensor sharing the dlpack data.
Parameters
----------
dlpack : PyCapsule
The capsule object of a dlpack tensor.
Returns
-------
dragon.EagerTensor
The tensor with the dlpack data.
"""
current_ws = workspace.get_workspace()
tensor = EagerTensor(device=None)
tensor._gc = current_ws.collectors.TENSOR
tensor._impl = current_ws.create_tensor(
tensor._gc.alloc('${DLPACK}')).FromDLPack(dlpack)
tensor._device = device_spec.DeviceSpec(*tensor._impl.device)
return tensor
def shape_args(**kwargs):
return {'device': device_spec.DeviceSpec('cpu')}
def new_device(device_type, device_index):
"""Return a new device abstraction."""
return device_spec.DeviceSpec(device_type, device_index)
def __init__(self, key, dev, **kwargs):
super(Shape, self).__init__(key, dev, **kwargs)
self._device = device_spec.DeviceSpec()