Fix reading FP16 scalars from TensorFlow (#4935)

This commit is contained in:
Maxim Vafin 2021-03-25 15:59:02 +03:00 committed by GitHub
parent fa37277e3e
commit 2a52747b03
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 3 deletions

View File

@ -17,6 +17,8 @@
import numpy as np
from tensorflow.core.framework import types_pb2 as tf_types # pylint: disable=no-name-in-module,import-error
# Suppress false positive pylint warning about function with too many arguments
# pylint: disable=E1121
# mapping between TF data type and numpy data type and function to extract data from TF tensor
_tf_np_mapping = [('DT_BOOL', np.bool, lambda pb: pb.bool_val, lambda x: bool_cast(x)),
('DT_INT8', np.int8, lambda pb: pb.int_val, lambda x: np.int8(x)),
@ -27,16 +29,16 @@ _tf_np_mapping = [('DT_BOOL', np.bool, lambda pb: pb.bool_val, lambda x: bool_ca
('DT_UINT16', np.uint16, lambda pb: pb.int_val, lambda x: np.uint16(x)),
('DT_UINT32', np.uint32, lambda pb: pb.uint32_val, lambda x: np.uint32(x)),
('DT_UINT64', np.uint64, lambda pb: pb.uint64_val, lambda x: np.uint64(x)),
('DT_HALF', np.float16, lambda pb: pb.half_val, lambda x: np.float16(x)),
('DT_HALF', np.float16, lambda pb: np.uint16(pb.half_val).view(np.float16), lambda x: np.float16(x)),
('DT_FLOAT', np.float32, lambda pb: pb.float_val, lambda x: np.float32(x)),
('DT_DOUBLE', np.double, lambda pb: pb.double_val, lambda x: np.double(x)),
('DT_STRING', np.str, lambda pb: pb.string_val, lambda x: np.str(x)),
]
tf_data_type_decode = {getattr(tf_types, tf_dt): (np_type, func) for tf_dt, np_type, func, cast in _tf_np_mapping if
tf_data_type_decode = {getattr(tf_types, tf_dt): (np_type, func) for tf_dt, np_type, func, _ in _tf_np_mapping if
hasattr(tf_types, tf_dt)}
tf_data_type_cast = {np_type: cast for tf_dt, np_type, func, cast in _tf_np_mapping if hasattr(tf_types, tf_dt)}
tf_data_type_cast = {np_type: cast for tf_dt, np_type, _, cast in _tf_np_mapping if hasattr(tf_types, tf_dt)}
def bool_cast(x):