def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:
"""Converts a tensor def object to a numpy array.
Inputs:
tensor: a TensorProto object.
base_dir: if external tensor exists, base_dir can help to find the path to it
Returns:
arr: the converted array.
"""
if tensor.HasField("segment"):
raise ValueError("Currently not supporting loading segments.")
if tensor.data_type == TensorProto.UNDEFINED:
raise TypeError("The element type in the input tensor is not defined.")
tensor_dtype = tensor.data_type
np_dtype = mapping.TENSOR_TYPE_TO_NP_TYPE[tensor_dtype]
storage_type = mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[tensor_dtype]
storage_np_dtype = mapping.TENSOR_TYPE_TO_NP_TYPE[storage_type]
storage_field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[storage_type]
dims = tensor.dims
if tensor.data_type == TensorProto.STRING:
utf8_strings = getattr(tensor, storage_field)
ss = list(s.decode('utf-8') for s in utf8_strings)
return np.asarray(ss).astype(np_dtype).reshape(dims)
# Load raw data from external tensor if it exists
if uses_external_data(tensor):
load_external_data_for_tensor(tensor, base_dir)
if tensor.HasField("raw_data"):
# Raw_bytes support: using frombuffer.
if sys.byteorder == 'big':
# Convert endian from little to big
convert_endian(tensor)
# manually convert bf16 since there's no numpy support
if tensor_dtype == TensorProto.BFLOAT16:
data = np.frombuffer(tensor.raw_data, dtype=np.int16)
return bfloat16_to_float32(data, dims)
return np.frombuffer(tensor.raw_data, dtype=np_dtype).reshape(dims)
else:
# float16 is stored as int32 (uint16 type); Need view to get the original value
if tensor_dtype == TensorProto.FLOAT16:
return (np.asarray(tensor.int32_data,
dtype=np.uint16).reshape(dims).view(np.float16))
# bfloat16 is stored as int32 (uint16 type); no numpy support for bf16
if tensor_dtype == TensorProto.BFLOAT16:
data = np.asarray(tensor.int32_data, dtype=np.int32)
return bfloat16_to_float32(data, dims)
data = getattr(tensor, storage_field)
if (tensor_dtype == TensorProto.COMPLEX64
or tensor_dtype == TensorProto.COMPLEX128):
data = combine_pairs_to_complex(data)
return (np.asarray(
data, dtype=storage_np_dtype).astype(np_dtype).reshape(dims))
def test_save_model_with_existing_raw_data_should_override(self) -> None:
model_file_path = self.get_temp_model_filename()
original_raw_data = self.model.graph.initializer[0].raw_data
onnx.save_model(self.model,
model_file_path,
save_as_external_data=True,
size_threshold=0)
self.assertTrue(Path.isfile(model_file_path))
model = onnx.load_model(model_file_path, load_external_data=False)
initializer_tensor = model.graph.initializer[0]
initializer_tensor.raw_data = b'dummpy_raw_data'
# If raw_data and external tensor exist at the same time, override existing raw_data
load_external_data_for_tensor(initializer_tensor, self.temp_dir)
self.assertEqual(initializer_tensor.raw_data, original_raw_data)