def MakeDeviceOption(device_type, gpu_id, rng_seed=None):
""" return a DeviceOption """
option = pb.DeviceOption()
option.device_type = device_type
option.gpu_id = gpu_id
if rng_seed is not None: option.random_seed = rng_seed
return option
def GraphDef_Device(meta_graph):
"""Inject the device option into GraphDef.
Parameters
----------
meta_graph : dragon_pb2.GraphDef
The definition of meta graph.
Returns
-------
None
References
----------
`config.EnableCPU()`_ - How to use CPU device.
`config.EnableCUDA(*args, **kwargs)`_ - How to use CUDA device.
`config.SetRandomSeed(*args, **kwargs)`_ - How to set random seed.
"""
from dragon.config import option
if option['device'] is not 'None':
supports = {'CPU': 0, 'CUDA': 1}
device_option = pb.DeviceOption()
device_option.device_type = supports[option['device']]
device_option.gpu_id = option['gpu_id']
device_option.random_seed = option['random_seed']
if option['use_cudnn']: device_option.engine = 'CUDNN'
meta_graph.device_option.CopyFrom(device_option)
def MakeDeviceOption(device_type, device_id, engine=None, rng_seed=None):
option = pb.DeviceOption()
option.device_type = device_type
option.device_id = device_id
if engine is not None: option.engine = engine
if rng_seed is not None: option.random_seed = rng_seed
return option
def FeedTensor(tensor, ndarray, force_cpu=False, dtype=None):
tensor = tensor.name if hasattr(tensor, 'name') else str(tensor)
dev = None
if force_cpu is True: dev = utils.MakeDeviceOption(0, 0)
else:
from dragon.core.scope import DEVICE_SCOPE
if DEVICE_SCOPE != '':
supports = {'/cpu': 0, '/gpu': 1}
dev = pb.DeviceOption()
dev.device_type = supports[DEVICE_SCOPE.split(':')[0]]
dev.gpu_id = int(DEVICE_SCOPE.split(':')[1])
else:
from dragon.config import option
if option['device'] == 'CUDA':
dev = utils.MakeDeviceOption(1, option['gpu_id'])
elif option['device'] == 'CPU':
dev = utils.MakeDeviceOption(0, 0)
if not isinstance(ndarray, np.ndarray):
if not isinstance(ndarray, list):
ndarray = [ndarray]
dtype = np.float32 if dtype is None else dtype
else:
dtype = ndarray.dtype if dtype is None else dtype
ndarray = np.array(ndarray, dtype=dtype)
FeedTensorCC(tensor, ndarray, StringfyProto(dev))
def GraphDef_Device(graph_def):
""" generate deivce info for CC Graph """
from dragon.config import option
if option['device'] is not 'None':
supports = {'CPU': 0, 'CUDA': 1}
device_option = pb.DeviceOption()
device_option.device_type = supports[option['device']]
device_option.gpu_id = option['gpu_id']
device_option.random_seed = option['random_seed']
if option['use_cudnn']: device_option.engine = 'CUDNN'
graph_def.device_option.CopyFrom(device_option)
def CreateOperator(cls, inputs, op_type, nout=None, existing_outputs=None,
extra_inputs=None, name=None, **kwargs):
expressions = OrderedDict() # keep order for displaying
# 1. collect inputs
if not isinstance(inputs, list): inputs = [inputs]
for input in inputs:
for op_idx, expr in input.expressions.items():
if not op_idx in expressions:
expressions[op_idx] = expr
if extra_inputs is not None:
if not isinstance(extra_inputs, list): extra_inputs = [extra_inputs]
for input in extra_inputs:
for op_idx, expr in input.expressions.items():
if not op_idx in expressions:
expressions[op_idx] = expr
# 2. generate outputs
outputs = []
if existing_outputs is None:
for idx in xrange(nout): outputs.append(Tensor())
else:
if not isinstance(existing_outputs, list): existing_outputs = [existing_outputs]
outputs = existing_outputs
nout = len(outputs)
if not isinstance(outputs, list): outputs = [outputs]
# 3. make def, then push back to expressions
inputs_name = [input.name for input in inputs]
outputs_name = [output.name for output in outputs]
op_idx, op_name = GetOperatorName(name)
device_option = None
from dragon.core.scope import DEVICE_SCOPE, ENGINE_SCOPE
if DEVICE_SCOPE != '':
supports = {'/cpu': 0, '/gpu': 1}
device_option = pb.DeviceOption()
device_option.device_type = supports[DEVICE_SCOPE.split(':')[0]]
device_option.gpu_id = int(DEVICE_SCOPE.split(':')[1])
device_option.engine = ENGINE_SCOPE
op_def = MakeOperatorDef(op_type, inputs_name, outputs_name, op_name,
device_option=device_option, **kwargs)
expressions[op_idx] = op_def
# 4. deliver expression & extra_targets to all outputs
for output in outputs:
output.expressions = expressions
# deliver extra_targets if necessary
for input in inputs:
output.extra_targets = \
output.extra_targets.union(input.extra_targets)
if extra_inputs is not None:
for input in extra_inputs:
output.extra_targets.add(input.name)
# 5. utils
if 'static_shape' in kwargs:
outputs[0].tf_shape = kwargs['static_shape']
if nout > 1:
return outputs
elif nout == 1:
return outputs[0]
else:
return None
def CreateOperator(cls,
inputs,
op_type,
nout=None,
existing_outputs=None,
output_shapes=None,
extra_inputs=None,
name=None,
**kwargs):
"""Construct a new Tensor with specific operator descriptor.
Parameters
----------
inputs : list of Tensor or Tensor
The inputs for this operator.
op_type : str
The operator type.
nout : int
The number of outputs to return.
It will be discarded if ``existing_outputs`` is not None.
existing_outputs : list of Tensor, Tensor or None
The existing outputs for this operator.
extra_inputs : list of Tensor, Tensor or None
The inputs that should be attached to solving targets, e.g. dynamic shape.
name : str or None
The optional name to use. ``Op_xxx`` will be used automatically if it is None.
Returns
-------
list of Tensor, Tensor, or None
The outputs of this operator.
Examples
--------
>>> a = Tensor().Variable()
>>> b = Tensor().Variable()
>>> c = Tensor.CreateOperator(inputs=[a, b], op_type='Add', nout=1)
>>> a = Tensor().Variable()
>>> b = Tensor().Variable()
>>> c = Tensor().Variable()
>>> c = Tensor.CreateOperator(inputs=[a, b], op_type='Add', existing_outputs=c)
>>> import dragon.core.workspace as ws
>>> import dragon.vm.theano as theano
>>> dynamic_shape = Tensor().Variable()
>>> ws.FeedTensor(dynamic_shape, [1, 2, 3, 4])
>>> a = ops.Fill(shape=dynamic_shape, value=5.0)
>>> print theano.function(outputs=a)
>>> [[ 5. 5. 5.]
[ 5. 5. 5.]]
"""
expressions = OrderedDict() # keep order for displaying
# 1. collect inputs
if not isinstance(inputs, list): inputs = [inputs]
for input in inputs:
for op_idx, expr in input.expressions.items():
if not op_idx in expressions:
expressions[op_idx] = expr
if extra_inputs is not None:
if not isinstance(extra_inputs, list):
extra_inputs = [extra_inputs]
for input in extra_inputs:
for op_idx, expr in input.expressions.items():
if not op_idx in expressions:
expressions[op_idx] = expr
# 2. generate outputs
outputs = []
if existing_outputs is None:
for idx in range(nout):
outputs.append(Tensor())
else:
if not isinstance(existing_outputs, list):
existing_outputs = [existing_outputs]
outputs = existing_outputs
nout = len(outputs)
if not isinstance(outputs, list): outputs = [outputs]
# 3. make def, then push expressions
inputs_name = [input.name for input in inputs]
outputs_name = [output.name for output in outputs]
op_idx, op_name = GetOperatorName(name)
device_option = None
from dragon.core.scope import _DEVICE_SCOPE, _ENGINE_SCOPE
if _DEVICE_SCOPE != '':
supports = {'/cpu': 0, '/gpu': 1}
device_option = pb.DeviceOption()
device_option.device_type = supports[_DEVICE_SCOPE.split(':')[0]]
device_option.gpu_id = int(_DEVICE_SCOPE.split(':')[1])
device_option.engine = _ENGINE_SCOPE
op_def = MakeOperatorDef(op_type,
inputs_name,
outputs_name,
op_name,
device_option=device_option,
**kwargs)
expressions[op_idx] = op_def
# 4. make outputs
for idx, output in enumerate(outputs):
# deliver expressions
output.expressions = expressions
# deliver extra targets
for input in inputs:
output.extra_targets = \
output.extra_targets.union(input.extra_targets)
if extra_inputs is not None:
for input in extra_inputs:
output.extra_targets.add(input.name)
# 5. utils
if 'static_shape' in kwargs:
outputs[0].tf_shape = kwargs['static_shape']
# 6. returns
if nout > 1: return outputs
elif nout == 1: return outputs[0]
else: return None
def FeedTensor(tensor, array, force_cpu=False, dtype=None):
"""Feed the values to the given tensor.
Parameters
----------
tensor : Tensor or str
The tensor to feed.
ndarray : number, list or ndarray
The values to feed.
force_cpu : boolean
Whether force to feed to cpu context.
dtype : str
The data type. If ``None``, ``float32`` will be used instead.
Returns
-------
None
Examples
--------
>>> import dragon as dg
>>> a = dg.Tensor().Variable()
>>> dg.workspace.FeedTensor(a, 1)
>>> a_value = dg.workspace.FetchTensor(a)
>>> a_value, a_value.dtype
>>> [ 1.], float32
>>> dg.workspace.FeedTensor(a, [[1, 2, 3]], dtype='float16')
>>> a_value = a.get_value()
>>> a_value, a_value.dtype
>>> [[ 1. 2. 3.]], float16
References
----------
The wrapper of ``FeedTensorCC``.
"""
name = tensor.name if hasattr(tensor, 'name') else str(tensor)
if force_cpu is True:
dev = utils.MakeDeviceOption(0, 0)
else:
from dragon.core.scope import _DEVICE_SCOPE
if _DEVICE_SCOPE != '':
supports = {'/cpu': 0, '/gpu': 1, '/mlu': 2}
dev = pb.DeviceOption()
dev.device_type = supports[_DEVICE_SCOPE.split(':')[0]]
dev.device_id = int(_DEVICE_SCOPE.split(':')[1])
else:
from dragon.config import option
if option['device'] == 'CUDA':
dev = utils.MakeDeviceOption(1, option['device_id'])
else:
dev = utils.MakeDeviceOption(0, 0)
if not isinstance(array, np.ndarray):
if not isinstance(array, list):
array = [array]
auto_data_type = np.float32 if dtype is None else dtype
else:
auto_data_type = array.dtype if dtype is None else dtype
if hasattr(tensor, 'dtype') and tensor.dtype is not None:
if tensor.dtype not in _DATA_TYPES:
raise TypeError('Unsupported data types: {}.'.format(tensor.dtype))
preset_data_type = _DATA_TYPES[tensor.dtype]
if dtype is not None:
if dtype != preset_data_type:
raise TypeError(
'The preset data type is {}, but force to {}.'.format(
preset_data_type, dtype))
auto_data_type = preset_data_type
nd_array = np.array(array, dtype=auto_data_type, copy=False)
FeedTensorCC(name, nd_array, _stringify_proto(dev))
