Repository: sql-machine-learning/elasticdl

Issue:
Transform BytesCodec into a general codec.
now the BytesCodec is for mnist dataset specified, and it should be modified to encode/decode any dataset using bytes format.


Discussion:

Resolving patch:
diff --git a/record_codec/__init__.py b/record_codec/__init__.py
--- a/record_codec/__init__.py
+++ b/record_codec/__init__.py
@@ -1,3 +1,4 @@
+from record_codec.bytes_codec import BytesCodec
 from record_codec.tf_example_codec import TFExampleCodec
 
-__all__=[TFExampleCodec]
+__all__=[BytesCodec, TFExampleCodec]
diff --git a/record_codec/bytes_codec.py b/record_codec/bytes_codec.py
--- a/record_codec/bytes_codec.py
+++ b/record_codec/bytes_codec.py
@@ -1,32 +1,42 @@
 import numpy as np
 
-N = 28
 
-
-# TODO: maybe use TF variant tensor to do more flexible encoding.
 class BytesCodec(object):
+    def __init__(self, feature_columns):
+        self._feature_columns = feature_columns
+        self._col_id = {
+            c.name: order for order, c in enumerate(feature_columns)
+        }
+
     def encode(self, data):
-        values = [] 
-        for _, f_value in data:
-            values.append(f_value)
-        return np.concatenate(values, axis=None).tobytes()
+        # Rearrange the data in order of the columns.
+        values = [None] * len(self._feature_columns)
+        for f_name, f_value in data:
+            col_id = self._col_id[f_name]
+            column = self._feature_columns[col_id]
+            if column.dtype != f_value.dtype or column.shape != f_value.shape:
+                raise ValueError(
+                    "Input data doesn't match column %s definition: column: (%s, %s) data: (%s, %s)" % (
+                    f_name, column.dtype, column.shape, f_value.dtype, f_value.shape)
+                )
+            values[col_id] = f_value.tobytes()
+        for id, value in enumerate(values):
+            if value is None:
+                raise ValueError(
+                    "Missing value for column: %s",
+                    self._col_id[id].name
+                )
+        return b"".join(values)
 
     def decode(self, record):
-        parsed = np.frombuffer(record, dtype="uint8")
-        assert len(parsed) == N * N + 1
-        label = parsed[-1]
-        parsed = np.resize(parsed[:-1], new_shape=(N, N))
-        return {'image': parsed, 'label': label}
-
-    def show(self, data, label):
-        """Print the image and label on terminal for debugging"""
-        assert data.shape == (N, N) and data.dtype == "uint8"
-        assert label >= 0 and label <= 9
-
-        def grey(x):
-            return "\033[48;2;%d;%d;%dm" % (x, x, x) + " \033[0m"
-
-        for line in data:
-            s = "".join(grey(x) for x in line)
-            print(s)
-        print("label =", label)
+        offset = 0
+        res = {}
+        for c in self._feature_columns:
+            count = np.prod(c.shape)
+            res[c.name] = np.frombuffer(
+                record,
+                dtype=c.dtype.as_numpy_dtype,
+                count=count,
+                offset=offset).reshape(c.shape)
+            offset += count * c.dtype.size
+        return res


Test patch:
diff --git a/record_codec/bytes_codec_test.py b/record_codec/bytes_codec_test.py
--- a/record_codec/bytes_codec_test.py
+++ b/record_codec/bytes_codec_test.py
@@ -1,18 +1,47 @@
 import unittest
-from elasticdl.record_codec.bytes_codec import BytesCodec 
+import tempfile
+import os
+import tensorflow as tf
 import numpy as np
+from recordio import File
+from record_codec.bytes_codec import BytesCodec
 
 
-class RecordTest(unittest.TestCase):
-    def test_round_trip(self):
-        # a random array
-        codec = BytesCodec()
-        d = np.ndarray(shape=(28, 28), dtype="uint8")
-        data = [('image', d), ('label', 5)]
-        encoded = codec.encode(data)
-        decoded = codec.decode(encoded)
-        np.testing.assert_array_equal(d, decoded['image'])
-        self.assertEqual(label, decoded['label'])
+class TestBytesCodec(unittest.TestCase):
+    """ Test bytes_codec.py
+    """
+
+    def test_encode_and_decode(self):
+        feature_columns = [tf.feature_column.numeric_column(key="f0",
+            dtype=tf.float64, shape=[1]),
+            tf.feature_column.numeric_column(key="label",
+            dtype=tf.int64, shape=[1])]
+
+        example_1 = [("f0", np.array([100.1])), ("label", np.array([1]))]
+        example_2 = [("f0", np.array([200.1])), ("label", np.array([2]))]
+        example_3 = [("f0", np.array([300.1])), ("label", np.array([3]))]
+        examples = [example_1, example_2, example_3]
+
+        tmp_file = tempfile.NamedTemporaryFile(delete=False)
+
+        # Create the codec for tf.train.Exampel data.
+        codec = BytesCodec(feature_columns)
+
+        # Write an encoded RecordIO file.
+        with File(tmp_file.name, "w", encoder=codec.encode) as coded_w:
+            for example in examples:
+                coded_w.write(example)
+
+        # Verify decoded content, with get() interface.
+        with File(tmp_file.name, "r", decoder=codec.decode) as coded_r:
+            for idx in range(coded_r.count()):
+                exp = coded_r.get(idx)
+                expected_exp = examples[idx]
+                f_0, label = exp["f0"], exp["label"]
+                self.assertEqual(f_0, expected_exp[0][1])
+                self.assertEqual(label, expected_exp[1][1])
+
+        os.remove(tmp_file.name)
 
 
 if __name__ == "__main__":