Repository: dask/distributed

Issue:
Serialize Keras Model
For Dask.distributed to work well with a deep learning library like Keras, it may be helpful for Dask to know how to serialize Keras models.

It looks like there are `keras.models.Model.save_weights_to_hdf5_group` and `keras.models.load_model` functions that store/load a Keras model to and from an HDF5 file.  We can probably copy most of this logic to create Dask serialize/deserialize functions that convert a Model to a series of frames of bytes. See [dask serialization docs for detailed information](http://distributed.readthedocs.io/en/latest/protocol.html#serializing-data).

Additionally, we would want to extend the `sizeof` function in `distributed.sizeof` to quickly provide a rough estimate of how many bytes are contained within a model (hopefully without having to serialize everything).

If we have these two pieces then Keras models should be able to move freely throughout Dask clusters.

This issue does require skill and attention to detail, but fortunately does not require much expertise in Dask itself.


Discussion:
cc @seibert just keeping you in the loop here.

Resolving patch:
diff --git a/distributed/protocol/__init__.py b/distributed/protocol/__init__.py
--- a/distributed/protocol/__init__.py
+++ b/distributed/protocol/__init__.py
@@ -21,3 +21,7 @@ def _register_h5py():
 @partial(register_serialization_lazy, "netCDF4")
 def _register_netcdf4():
     from . import netcdf4
+
+@partial(register_serialization_lazy, "keras")
+def _register_keras():
+    from . import keras
diff --git a/distributed/protocol/keras.py b/distributed/protocol/keras.py
new file mode 100644
--- /dev/null
+++ b/distributed/protocol/keras.py
@@ -0,0 +1,41 @@
+from __future__ import print_function, division, absolute_import
+
+from .serialize import register_serialization, serialize, deserialize
+
+from ..utils import log_errors, ensure_bytes
+
+
+def serialize_keras_model(model):
+    import keras
+    if keras.__version__ < '1.2.0':
+        raise ImportError("Need Keras >= 1.2.0. "
+                "Try pip install keras --upgrade --no-deps")
+
+    header = model._updated_config()
+    weights = model.get_weights()
+    headers, frames = list(zip(*map(serialize, weights)))
+    header['headers'] = headers
+    header['nframes'] = [len(L) for L in frames]
+    frames = [frame for L in frames for frame in L]
+    return header, frames
+
+
+def deserialize_keras_model(header, frames):
+    from keras.models import model_from_config
+    n = 0
+    weights = []
+    for head, length in zip(header['headers'], header['nframes']):
+        x = deserialize(head, frames[n: n + length])
+        weights.append(x)
+        n += length
+    model = model_from_config(header)
+    model.set_weights(weights)
+    return model
+
+
+register_serialization('keras.engine.training.Model', serialize_keras_model,
+                       deserialize_keras_model)
+register_serialization('keras.models.Model', serialize_keras_model,
+                       deserialize_keras_model)
+register_serialization('keras.models.Sequential', serialize_keras_model,
+                       deserialize_keras_model)
diff --git a/distributed/protocol/utils.py b/distributed/protocol/utils.py
--- a/distributed/protocol/utils.py
+++ b/distributed/protocol/utils.py
@@ -44,8 +44,7 @@ def merge_frames(header, frames):
     >>> merge_frames({'lengths': [6]}, [b'123', b'456'])
     [b'123456']
     """
-    lengths = list(header['lengths'])[::-1]
-    frames = list(frames)[::-1]
+    lengths = list(header['lengths'])
 
     if not frames:
         return frames
@@ -55,6 +54,9 @@ def merge_frames(header, frames):
     if all(len(f) == l for f, l in zip(frames, lengths)):
         return frames
 
+    frames = frames[::-1]
+    lengths = lengths[::-1]
+
     out = []
     while lengths:
         l = lengths.pop()


Test patch:
diff --git a/distributed/protocol/tests/test_keras.py b/distributed/protocol/tests/test_keras.py
new file mode 100644
--- /dev/null
+++ b/distributed/protocol/tests/test_keras.py
@@ -0,0 +1,29 @@
+import functools
+import sys
+import traceback
+
+import numpy as np
+from numpy.testing import assert_allclose
+import pytest
+
+keras = pytest.importorskip('keras')
+
+from distributed.protocol import serialize, deserialize, dumps, loads, to_serialize
+
+
+def test_serialize_deserialize_model():
+    model = keras.models.Sequential()
+    model.add(keras.layers.Dense(5, input_dim=3))
+    model.add(keras.layers.Dense(2))
+    model.compile(optimizer='sgd', loss='mse')
+    x = np.random.random((1, 3))
+    y = np.random.random((1, 2))
+    model.train_on_batch(x, y)
+
+    loaded = deserialize(*serialize(model))
+    assert_allclose(loaded.predict(x), model.predict(x))
+
+    data = {'model': to_serialize(model)}
+    frames = dumps(data)
+    result = loads(frames)
+    assert_allclose(result['model'].predict(x), model.predict(x))
diff --git a/distributed/protocol/tests/test_protocol_utils.py b/distributed/protocol/tests/test_protocol_utils.py
--- a/distributed/protocol/tests/test_protocol_utils.py
+++ b/distributed/protocol/tests/test_protocol_utils.py
@@ -11,3 +11,6 @@ def test_merge_frames():
 
     b = b'123'
     assert merge_frames({'lengths': [3]}, [b])[0] is b
+
+    L = [b'123', b'456']
+    assert merge_frames({'lengths': [3, 3]}, L) is L