Merge pull request PaddlePaddle#77 from jacobbieker/master

Add autencoding of numpy arrays

Author: tmbdev

webdataset/__init__.py

@@ -29,7 +29,7 @@
     ignore_and_stop,
     warn_and_stop,
 )
-from .writer import ShardWriter, TarWriter, torch_dumps
+from .writer import ShardWriter, TarWriter, torch_dumps, numpy_dumps
 from .autodecode import (
     Continue,
     handle_extension,

webdataset/writer.py

@@ -78,6 +78,19 @@ def torch_dumps(data):
     return stream.getvalue()
 
 
+def numpy_dumps(data):
+    """Dump data into a bytestring using numpy npy format
+    
+    :param data: data to be dumped
+    """
+    import io
+    import numpy.lib.format
+    
+    stream = io.BytesIO()
+    numpy.lib.format.write_array(stream, data)
+    return stream.getvalue()
+
+
 def make_handlers():
     """Create a list of handlers for encoding data."""
     handlers = {}
@@ -99,6 +112,8 @@ def f(extension_):
         handlers[extension] = pickle.dumps
     for extension in ["pth"]:
         handlers[extension] = torch_dumps
+    for extension in ["npy"]:
+        handlers[extension] = numpy_dumps
     for extension in ["json", "jsn"]:
         handlers[extension] = lambda x: json.dumps(x).encode("utf-8")
     for extension in ["ten", "tb"]: