Fix reading FP16 scalars from TensorFlow (#4935)
This commit is contained in:
parent
fa37277e3e
commit
2a52747b03
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in New Issue