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)