def make_tensor(name, data_type, dims, vals, raw=False):
'''
Make a TensorProto with specified arguments. If raw is False, this
function will choose the corresponding proto field to store the
values based on data_type. If raw is True, use "raw_data" proto
field to store the values, and values should be of type bytes in
this case.
'''
tensor = TensorProto()
tensor.data_type = data_type
tensor.name = name
if data_type == TensorProto.STRING:
assert not raw, "Can not use raw_data to store string type"
tensor.string_data.extend(vals)
if (data_type == TensorProto.COMPLEX64
or data_type == TensorProto.COMPLEX128):
vals = split_complex_to_pairs(vals)
if raw:
tensor.raw_data = vals
else:
field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[
mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[data_type]]
getattr(tensor, field).extend(vals)
tensor.dims.extend(dims)
return tensor
def make_tensor(name, data_type, dims, vals, raw=False):
'''
Make a TensorProto with specified arguments. If raw is False, this
function will choose the corresponding proto field to store the
values based on data_type. If raw is True, use "raw_data" proto
field to store the values, and values should be of type bytes in
this case.
'''
tensor = TensorProto()
tensor.data_type = data_type
tensor.name = name
if data_type == TensorProto.STRING:
assert not raw, "Can not use raw_data to store string type"
tensor.string_data.extend(vals)
elif data_type in [
TensorProto.UINT8, TensorProto.INT8, TensorProto.UINT16,
TensorProto.INT16, TensorProto.INT32, TensorProto.FLOAT16,
TensorProto.BOOL, TensorProto.FLOAT
]:
if raw:
tensor.raw_data = vals
else:
if data_type == TensorProto.FLOAT:
tensor.float_data.extend(vals)
elif data_type == TensorProto.INT64:
tensor.int64_data.extend(vals)
else:
tensor.int32_data.extend(vals)
else:
raise RuntimeError('Unrecognized data_type: {}'.format(data_type))
tensor.dims.extend(dims)
return tensor
def from_array(arr, name=None):
"""Converts a numpy array to a tensor def.
Inputs:
arr: a numpy array.
name: (optional) the name of the tensor.
Returns:
tensor_def: the converted tensor def.
"""
tensor = TensorProto()
tensor.dims.extend(arr.shape)
if name:
tensor.name = name
if arr.dtype == np.object:
# Special care for strings.
raise NotImplementedError("Need to properly implement string.")
# For numerical types, directly use numpy raw bytes.
try:
dtype = mapping.NP_TYPE_TO_TENSOR_TYPE[arr.dtype]
except KeyError:
raise RuntimeError("Numpy data type not understood yet: {}".format(
str(arr.dtype)))
tensor.data_type = dtype
tensor.raw_data = arr.tobytes() # note: tobytes() is only after 1.9.
return tensor
def test_check_string_tensor(self):
tensor = TensorProto()
tensor.data_type = TensorProto.STRING
tensor.string_data.append('Test'.encode('utf-8'))
checker.check_tensor(tensor)
del tensor.string_data[:]
tensor.raw_data = 'Test'.encode('utf-8')
# string data should not be stored in raw_data field
self.assertRaises(ValueError, checker.check_tensor, tensor)
def from_array(arr, name=None):
"""Converts a numpy array to a tensor def.
Inputs:
arr: a numpy array.
name: (optional) the name of the tensor.
Returns:
tensor_def: the converted tensor def.
"""
tensor = TensorProto()
tensor.dims.extend(arr.shape)
if name:
tensor.name = name
dtype_map = {
np.dtype("float32"): TensorProto.FLOAT,
np.dtype("uint8"): TensorProto.UINT8,
np.dtype("int8"): TensorProto.INT8,
np.dtype("uint16"): TensorProto.UINT16,
np.dtype("int16"): TensorProto.INT16,
np.dtype("int32"): TensorProto.INT32,
np.dtype("int64"): TensorProto.INT64,
np.dtype("bool"): TensorProto.BOOL,
np.dtype("float16"): TensorProto.FLOAT16,
}
if arr.dtype == np.object:
# Special care for strings.
raise NotImplementedError("Need to properly implement string.")
# For numerical types, directly use numpy raw bytes.
try:
dtype = dtype_map[arr.dtype]
except KeyError:
raise RuntimeError(
"Numpy data type not understood yet: {}".format(str(arr.dtype)))
tensor.data_type = dtype
tensor.raw_data = arr.tobytes() # note: tobytes() is only after 1.9.
return tensor