Repository: sql-machine-learning/elasticdl

Issue:
Support Arbitrary Data type in protobuf
Currently, we only support data type float32 in protobuf. 

I am planning to refactor the evaluation process to fix issue #1082. To support metrics inherited from `tf.Metric`, worker need to transport labels and model outputs to master. We can't suppose that labels and model outputs are of type `float32`. So we need to support arbitrary data type in protobuf.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/dtypes.py b/elasticdl/python/common/dtypes.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/common/dtypes.py
@@ -0,0 +1,40 @@
+import numpy as np
+
+from elasticdl.proto import tensor_dtype_pb2
+
+
+def dtype_tensor_to_numpy(dtype):
+    """Convert tensor dtype to numpy dtype object."""
+    return np.dtype(_DT_TENSOR_TO_NP.get(dtype, None))
+
+
+def dtype_numpy_to_tensor(dtype):
+    """Convert numpy dtype object to tensor dtype."""
+    return _DT_NP_TO_TENSOR.get(dtype.type, tensor_dtype_pb2.DT_INVALID)
+
+
+def is_numpy_dtype_allowed(dtype):
+    return dtype.type in _DT_NP_TO_TENSOR
+
+
+_DT_TENSOR_TO_NP = {
+    tensor_dtype_pb2.DT_INT8: np.int8,
+    tensor_dtype_pb2.DT_INT16: np.int16,
+    tensor_dtype_pb2.DT_INT32: np.int32,
+    tensor_dtype_pb2.DT_INT64: np.int64,
+    tensor_dtype_pb2.DT_FLOAT16: np.float16,
+    tensor_dtype_pb2.DT_FLOAT32: np.float32,
+    tensor_dtype_pb2.DT_FLOAT64: np.float64,
+    tensor_dtype_pb2.DT_BOOL: np.bool,
+}
+
+_DT_NP_TO_TENSOR = {
+    np.int8: tensor_dtype_pb2.DT_INT8,
+    np.int16: tensor_dtype_pb2.DT_INT16,
+    np.int32: tensor_dtype_pb2.DT_INT32,
+    np.int64: tensor_dtype_pb2.DT_INT64,
+    np.float16: tensor_dtype_pb2.DT_FLOAT16,
+    np.float32: tensor_dtype_pb2.DT_FLOAT32,
+    np.float64: tensor_dtype_pb2.DT_FLOAT64,
+    np.bool: tensor_dtype_pb2.DT_BOOL,
+}
diff --git a/elasticdl/python/common/ndarray.py b/elasticdl/python/common/ndarray.py
--- a/elasticdl/python/common/ndarray.py
+++ b/elasticdl/python/common/ndarray.py
@@ -1,7 +1,12 @@
 import numpy as np
 import tensorflow as tf
 
-from elasticdl.proto import elasticdl_pb2
+from elasticdl.proto import elasticdl_pb2, tensor_dtype_pb2
+from elasticdl.python.common.dtypes import (
+    dtype_numpy_to_tensor,
+    dtype_tensor_to_numpy,
+    is_numpy_dtype_allowed,
+)
 
 
 def tensor_to_ndarray(tensor_pb):
@@ -14,8 +19,14 @@ def tensor_to_ndarray(tensor_pb):
     if not tensor_pb.dim:
         raise ValueError("Tensor PB has no dim defined")
 
+    # TODO (yunjian.lmh): Tensor of checkpoint file in `tests/testdata` does
+    # not have `dtype` field. Remove default value for `tensor_pb.dtype` after
+    # generating a new checkpoint file.
+    if not tensor_pb.dtype:
+        tensor_pb.dtype = tensor_dtype_pb2.DT_FLOAT32
+    dtype = dtype_tensor_to_numpy(tensor_pb.dtype)
     # Check that the buffer size agrees with dimensions.
-    size = 4  # A float32 item occupies 4 bytes
+    size = dtype.itemsize
     for d in tensor_pb.dim:
         size *= d
     if size != len(tensor_pb.content):
@@ -26,11 +37,11 @@ def tensor_to_ndarray(tensor_pb):
         )
     if not tensor_pb.indices:
         arr = np.ndarray(
-            shape=tensor_pb.dim, dtype=np.float32, buffer=tensor_pb.content
+            shape=tensor_pb.dim, dtype=dtype, buffer=tensor_pb.content
         )
     else:
         values = np.ndarray(
-            shape=tensor_pb.dim, dtype=np.float32, buffer=tensor_pb.content
+            shape=tensor_pb.dim, dtype=dtype, buffer=tensor_pb.content
         )
         indices = tf.convert_to_tensor(tensor_pb.indices)
         arr = tf.IndexedSlices(values, indices)
@@ -42,14 +53,13 @@ def tensor_to_ndarray(tensor_pb):
 def ndarray_to_tensor(arr, indices=None):
     """Convert ndarray to Tensor PB"""
 
-    if arr.dtype != np.float32:
-        raise ValueError(
-            "expected ndarray to be of float32 type, got %s type", arr.dtype
-        )
+    if not is_numpy_dtype_allowed(arr.dtype):
+        raise ValueError("Dtype of ndarray %s is not supported", arr.dtype)
     tensor = elasticdl_pb2.Tensor()
     tensor.dim.extend(arr.shape)
     tensor.content = arr.tobytes()
     if indices:
         tensor.indices.extend(indices)
+    tensor.dtype = dtype_numpy_to_tensor(arr.dtype)
 
     return tensor


Test patch:
diff --git a/elasticdl/python/tests/ndarray_test.py b/elasticdl/python/tests/ndarray_test.py
--- a/elasticdl/python/tests/ndarray_test.py
+++ b/elasticdl/python/tests/ndarray_test.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from elasticdl.proto import elasticdl_pb2
+from elasticdl.proto import elasticdl_pb2, tensor_dtype_pb2
 from elasticdl.python.common.ndarray import (
     ndarray_to_tensor,
     tensor_to_ndarray,
@@ -12,7 +12,7 @@
 class ConverterTest(unittest.TestCase):
     def test_ndarray_to_tensor(self):
         # Wrong type, should raise
-        arr = np.array([1, 2, 3, 4])
+        arr = np.array([1, 2, 3, 4], dtype=np.uint8)
         self.assertRaises(ValueError, ndarray_to_tensor, arr)
 
         # Empty array
@@ -47,6 +47,7 @@ def testtensor_to_ndarray(self):
         # Empty array, should be ok.
         t.dim.append(0)
         t.content = b""
+        t.dtype = tensor_dtype_pb2.DT_FLOAT32
         arr = tensor_to_ndarray(t)
         np.testing.assert_array_equal(np.array([], dtype=np.float32), arr)
 
@@ -54,6 +55,7 @@ def testtensor_to_ndarray(self):
         del t.dim[:]
         t.dim.extend([2, 0, 1, 9])
         t.content = b""
+        t.dtype = tensor_dtype_pb2.DT_FLOAT32
         arr = tensor_to_ndarray(t)
         np.testing.assert_array_equal(
             np.ndarray(shape=[2, 0, 1, 9], dtype=np.float32), arr
@@ -64,6 +66,7 @@ def testtensor_to_ndarray(self):
         # Wrong content size, should raise
         del t.dim[:]
         t.dim.extend([11])
+        t.dtype = tensor_dtype_pb2.DT_FLOAT32
         self.assertRaises(ValueError, tensor_to_ndarray, t)
 
         # Compatible dimensions, should be ok.
@@ -71,6 +74,7 @@ def testtensor_to_ndarray(self):
             del t.dim[:]
             t.content = b"\0" * (4 * 12)
             t.dim.extend([m, 12 // m])
+            t.dtype = tensor_dtype_pb2.DT_FLOAT32
             arr = tensor_to_ndarray(t)
 
     def testRoundTrip(self):
@@ -78,11 +82,18 @@ def verify(a):
             b = tensor_to_ndarray(ndarray_to_tensor(a))
             np.testing.assert_array_equal(a, b)
 
+        # dtype = np.float32
         # 1-D array
         verify(np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32))
         # 4-D random array
         verify(np.ndarray(shape=[2, 1, 3, 4], dtype=np.float32))
 
+        # dtype = np.int64
+        # 1-D random array
+        verify(np.array([1, 2, 3, 4], dtype=np.int64))
+        # 4-D random array
+        verify(np.ndarray(shape=[2, 1, 3, 4], dtype=np.int64))
+
 
 if __name__ == "__main__":
     unittest.main()