diff --git a/onnx/helper.py b/onnx/helper.py index 8af88a70b20..21396aa3075 100644 --- a/onnx/helper.py +++ b/onnx/helper.py @@ -332,14 +332,22 @@ def make_tensor( if data_type == TensorProto.STRING: assert not raw, "Can not use raw_data to store string type" + np_dtype = mapping.TENSOR_TYPE_TO_NP_TYPE[data_type] + # Check number of vals specified equals tensor size - storage_type = mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[data_type] - expected_size = 1 if (not raw) else (mapping.TENSOR_TYPE_TO_NP_TYPE[storage_type].itemsize) - # Flatten a numpy array if its rank > 1 + expected_size = 1 + if raw: + # NumPy doesn't have BFLOAT16. TENSOR_TYPE_TO_NP_TYPE maps it to float32, + # which has the wrong itemsize. + if data_type == TensorProto.BFLOAT16: + expected_size = 2 + else: + expected_size = np_dtype.itemsize + if type(vals) is np.ndarray and len(vals.shape) > 1: vals = vals.flatten() for d in dims: - expected_size = expected_size * d + expected_size *= d if len(vals) != expected_size: raise ValueError("Number of values does not match tensor's size. Expected {}, but it is {}. " @@ -348,14 +356,12 @@ def make_tensor( if raw: tensor.raw_data = vals else: - if (data_type == TensorProto.COMPLEX64 - or data_type == TensorProto.COMPLEX128): + if (data_type == TensorProto.COMPLEX64 or data_type == TensorProto.COMPLEX128): vals = split_complex_to_pairs(vals) - # float16/bfloat16 are stored as uint16 elif data_type == TensorProto.FLOAT16: - vals = np.array(vals).astype(np.float16).view(dtype=np.uint16).flatten().tolist() + vals = np.array(vals).astype(np_dtype).view(dtype=np.uint16).flatten().tolist() elif data_type == TensorProto.BFLOAT16: - vals = list(map(float32_to_bfloat16, np.array(vals).astype(np.float32).flatten().tolist())) + vals = list(map(float32_to_bfloat16, np.array(vals).astype(np_dtype).flatten().tolist())) field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[ mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[data_type]] getattr(tensor, field).extend(vals) diff --git a/onnx/mapping.py b/onnx/mapping.py index f899b0547a8..be5b688fd74 100644 --- a/onnx/mapping.py +++ b/onnx/mapping.py @@ -27,7 +27,8 @@ # Numpy float32 array is only reversed to TensorProto.FLOAT NP_TYPE_TO_TENSOR_TYPE = {v: k for k, v in TENSOR_TYPE_TO_NP_TYPE.items() if k != TensorProto.BFLOAT16} -# This map indicates what storage-type is used in the protobuf (serialized) representation for TensorProto +# This is only used to get keys into STORAGE_TENSOR_TYPE_TO_FIELD. +# TODO(https://github.com/onnx/onnx/issues/4261): Remove this. TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE = { int(TensorProto.FLOAT): int(TensorProto.FLOAT), int(TensorProto.UINT8): int(TensorProto.INT32), diff --git a/onnx/test/helper_test.py b/onnx/test/helper_test.py index 3c8b88540b1..a0fcacad36f 100644 --- a/onnx/test/helper_test.py +++ b/onnx/test/helper_test.py @@ -392,6 +392,28 @@ def test_make_tensor(self) -> None: ) self.assertEqual(string_list, list(tensor.string_data)) + def test_make_int8_tensor(self) -> None: + np_array = np.random.randn(2, 3).astype(np.int8) + + tensor = helper.make_tensor( + name='test', + data_type=TensorProto.INT8, + dims=(2, 3), + vals=np_array + ) + self.assertEqual(tensor.name, 'test') + np.testing.assert_equal(np_array, numpy_helper.to_array(tensor)) + + # use raw_data field to store the data + tensor = helper.make_tensor( + name='test', + data_type=TensorProto.INT8, + dims=(2, 3), + vals=np_array.tobytes(), + raw=True, + ) + np.testing.assert_equal(np_array, numpy_helper.to_array(tensor)) + def test_make_float16_tensor(self) -> None: np_array = np.random.randn(2, 3).astype(np.float16) @@ -404,7 +426,7 @@ def test_make_float16_tensor(self) -> None: self.assertEqual(tensor.name, 'test') np.testing.assert_equal(np_array, numpy_helper.to_array(tensor)) - def test_make_float16_tensor_with_raw(self) -> None: + def test_make_float16_tensor_raw(self) -> None: np_array = np.random.randn(2, 3).astype(np.float16) tensor = helper.make_tensor( @@ -443,7 +465,7 @@ def test_make_bfloat16_tensor(self) -> None: self.assertEqual(tensor.name, 'test') np.testing.assert_equal(np_results, numpy_helper.to_array(tensor)) - def test_make_bfloat16_tensor_with_raw(self) -> None: + def test_make_bfloat16_tensor_raw(self) -> None: # numpy doesn't support bf16, so we have to compute the correct result manually np_array = np.array([[1.0, 2.0], [3.0, 4.0], [0.099853515625, 0.099365234375], [0.0998535081744, 0.1], [np.nan, np.inf]], dtype=np.float32)