Allow saving models directly to binary stream (keras-team#9789)

* Allow model saving/loading code to accept h5py.File objects.

This change allows saving models to memory-mapped files that do not have a physical presence on disk, and extraction of the serialized model data as a raw binary byte stream.

* Record file opening only after successfully opening the file.

keras/models.py

@@ -46,7 +46,9 @@ def save_model(model, filepath, overwrite=True, include_optimizer=True):
 
     # Arguments
         model: Keras model instance to be saved.
-        filepath: String, path where to save the model.
+        filepath: one of the following:
+            - string, path where to save the model, or
+            - h5py.File object where to save the model
         overwrite: Whether we should overwrite any existing
             model at the target location, or instead
             ask the user with a manual prompt.
@@ -97,13 +99,20 @@ def get_json_type(obj):
 
     from . import __version__ as keras_version
 
-    # If file exists and should not be overwritten.
-    if not overwrite and os.path.isfile(filepath):
-        proceed = ask_to_proceed_with_overwrite(filepath)
-        if not proceed:
-            return
+    if not isinstance(filepath, h5py.File):
+        # If file exists and should not be overwritten.
+        if not overwrite and os.path.isfile(filepath):
+            proceed = ask_to_proceed_with_overwrite(filepath)
+            if not proceed:
+                return
+
+        f = h5py.File(filepath, mode='w')
+        opened_new_file = True
+    else:
+        f = filepath
+        opened_new_file = False
 
-    with h5py.File(filepath, mode='w') as f:
+    try:
         f.attrs['keras_version'] = str(keras_version).encode('utf8')
         f.attrs['backend'] = K.backend().encode('utf8')
         f.attrs['model_config'] = json.dumps({
@@ -179,13 +188,18 @@ def get_json_type(obj):
                         else:
                             param_dset[:] = val
         f.flush()
+    finally:
+        if opened_new_file:
+            f.close()
 
 
 def load_model(filepath, custom_objects=None, compile=True):
     """Loads a model saved via `save_model`.
 
     # Arguments
-        filepath: String, path to the saved model.
+        filepath: one of the following:
+            - string, path to the saved model, or
+            - h5py.File object from which to load the model
         custom_objects: Optional dictionary mapping names
             (strings) to custom classes or functions to be
             considered during deserialization.
@@ -234,7 +248,14 @@ def convert_custom_objects(obj):
         if obj in custom_objects:
             return custom_objects[obj]
         return obj
-    with h5py.File(filepath, mode='r') as f:
+
+    opened_new_file = not isinstance(filepath, h5py.File)
+    if opened_new_file:
+        f = h5py.File(filepath, mode='r')
+    else:
+        f = filepath
+
+    try:
         # instantiate model
         model_config = f.attrs.get('model_config')
         if model_config is None:
@@ -292,7 +313,10 @@ def convert_custom_objects(obj):
                               'state. As a result, your model is '
                               'starting with a freshly initialized '
                               'optimizer.')
-    return model
+        return model
+    finally:
+        if opened_new_file:
+            f.close()
 
 
 def model_from_config(config, custom_objects=None):

tests/test_model_saving.py

@@ -104,6 +104,79 @@ def test_functional_model_saving():
     assert_allclose(out, out2, atol=1e-05)
 
 
+@keras_test
+def test_model_saving_to_pre_created_h5py_file():
+    inputs = Input(shape=(3,))
+    x = Dense(2)(inputs)
+    outputs = Dense(3)(x)
+
+    model = Model(inputs, outputs)
+    model.compile(loss=losses.MSE,
+                  optimizer=optimizers.Adam(),
+                  metrics=[metrics.categorical_accuracy])
+    x = np.random.random((1, 3))
+    y = np.random.random((1, 3))
+    model.train_on_batch(x, y)
+
+    out = model.predict(x)
+    _, fname = tempfile.mkstemp('.h5')
+    with h5py.File(fname, mode='r+') as h5file:
+        save_model(model, h5file)
+        loaded_model = load_model(h5file)
+        out2 = loaded_model.predict(x)
+    assert_allclose(out, out2, atol=1e-05)
+
+    # test non-default options in h5
+    with h5py.File('does not matter', driver='core',
+                   backing_store=False) as h5file:
+        save_model(model, h5file)
+        loaded_model = load_model(h5file)
+        out2 = loaded_model.predict(x)
+    assert_allclose(out, out2, atol=1e-05)
+
+
+@keras_test
+def test_model_saving_to_binary_stream():
+    inputs = Input(shape=(3,))
+    x = Dense(2)(inputs)
+    outputs = Dense(3)(x)
+
+    model = Model(inputs, outputs)
+    model.compile(loss=losses.MSE,
+                  optimizer=optimizers.Adam(),
+                  metrics=[metrics.categorical_accuracy])
+    x = np.random.random((1, 3))
+    y = np.random.random((1, 3))
+    model.train_on_batch(x, y)
+
+    out = model.predict(x)
+    _, fname = tempfile.mkstemp('.h5')
+    with h5py.File(fname, mode='r+') as h5file:
+        save_model(model, h5file)
+        loaded_model = load_model(h5file)
+        out2 = loaded_model.predict(x)
+    assert_allclose(out, out2, atol=1e-05)
+
+    # Save the model to an in-memory-only h5 file.
+    with h5py.File('does not matter', driver='core',
+                   backing_store=False) as h5file:
+        save_model(model, h5file)
+        h5file.flush()  # Very important! Otherwise you get all zeroes below.
+        binary_data = h5file.fid.get_file_image()
+
+        # Make sure the binary data is correct by saving it to a file manually
+        # and then loading it the usual way.
+        with open(fname, 'wb') as raw_file:
+            raw_file.write(binary_data)
+
+    # Load the manually-saved binary data, and make sure the model is intact.
+    with h5py.File(fname, mode='r') as h5file:
+        loaded_model = load_model(h5file)
+        out2 = loaded_model.predict(x)
+
+    assert_allclose(out, out2, atol=1e-05)
+
+
 @keras_test
 def test_saving_multiple_metrics_outputs():
     inputs = Input(shape=(5,))