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 _gen_op(self):
self._op = pb_utils.MakeOperatorDef(
op_type=self.op_meta['op_type'],
name='runtime',
inputs=[
'I({})'.format(i) for i in range(self.op_meta['n_inputs'])
],
outputs=[
'O({})'.format(i) for i in range(self.op_meta['n_outputs'])
],
device_option=pb_utils.MakeDeviceOption({
'CPU': 0,
'CUDA': 1
}[self._ctx[0]],
self._ctx[1],
engine='CUDNN'),
persistent_key=self.persistent_key,
**self.op_meta['arguments'])
dg.workspace.CreatePersistentOp(self._op)
return self._op
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))
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import dragon.core.utils as pb_utils
# Enumerate and persistent device options
# Make item frequently will degrade performance
DEVICE_LIMITS = 16
DEVICE_ENGINE = 'CUDNN'
CTX_TO_DEVICE_OPTION = {('CPU', 0): pb_utils.MakeDeviceOption(0, 0)}
for i in range(DEVICE_LIMITS):
CTX_TO_DEVICE_OPTION['CUDA', i] = \
pb_utils.MakeDeviceOption(1, i, DEVICE_ENGINE)
