tf keras savedmodel support

omegaml/backends/basemodel.py

@@ -28,7 +28,7 @@ def get(self, name, **kwargs):
         :param version: the version of the object (not supported)
         """
         # support new backend architecture while keeping back compatibility
-        return self.get_model(name)
+        return self.get_model(name, **kwargs)
 
     def put(self, obj, name, **kwargs):
         """
@@ -41,7 +41,7 @@ def put(self, obj, name, **kwargs):
         # support new backend architecture while keeping back compatibility
         return self.put_model(obj, name, **kwargs)
 
-    def put_model(self, obj, name, attributes=None):
+    def put_model(self, obj, name, attributes=None, **kwargs):
         """
         store a model
 
@@ -51,7 +51,7 @@ def put_model(self, obj, name, attributes=None):
         """
         raise NotImplementedError
 
-    def get_model(self, name, version=-1):
+    def get_model(self, name, version=-1, **kwargs):
         """
         retrieve a model
 

omegaml/backends/keras.py

@@ -26,7 +26,7 @@ def _load_model(self, fn):
         from keras.engine.saving import load_model
         return load_model(fn)
 
-    def put_model(self, obj, name, attributes=None):
+    def put_model(self, obj, name, attributes=None, **kwargs):
         fn = temp_filename()
         self._save_model(obj, fn)
         with open(fn, mode='rb') as fin:
@@ -85,8 +85,8 @@ def fit(self, modelname, Xname, Yname=None, validation_data=None,
                 valX = self.data_store.get(valX)
             if isinstance(Y, six.string_types):
                 valY = self.data_store.get(valY)
-            kwargs['validation_data'] = (valX, valY)
-        history = model.fit(X, Y, **kwargs)
+            keras_kwargs['validation_data'] = (valX, valY)
+        history = model.fit(X, Y, **keras_kwargs)
         meta = self.put_model(model, modelname, attributes={
             'history': serializable_history(history)
         })

omegaml/backends/tensorflow/__init__.py

@@ -1,3 +1,4 @@
 from .tfestimatormodel import TFEstimatorModelBackend, TFEstimatorModel
 from .tfkeras import TensorflowKerasBackend
 from .tfsavedmodel import TensorflowSavedModelPredictor, TensorflowSavedModelBackend
+from .tfkerassavedmodel import  TensorflowKerasSavedModelBackend

omegaml/backends/tensorflow/tfkeras.py

@@ -14,7 +14,7 @@ def supports(self, obj, name, **kwargs):
         import tensorflow as tf
         tfSequential = tf.keras.models.Sequential
         tfModel = tf.keras.models.Model
-        return isinstance(obj, (tfSequential, tfModel))
+        return isinstance(obj, (tfSequential, tfModel)) and not kwargs.get('as_savedmodel')
 
     def _save_model(self, model, fn):
         # override to implement model saving

omegaml/backends/tensorflow/tfkerassavedmodel.py

@@ -0,0 +1,28 @@
+import tempfile
+
+from omegaml.backends.tensorflow import TensorflowSavedModelBackend
+from omegaml.util import temp_filename
+
+
+class TensorflowKerasSavedModelBackend(TensorflowSavedModelBackend):
+    KIND = 'tfkeras.savedmodel'
+
+    @classmethod
+    def supports(self, obj, name, **kwargs):
+        import tensorflow as tf
+        tfSequential = tf.keras.models.Sequential
+        tfModel = tf.keras.models.Model
+        return isinstance(obj, (tfSequential, tfModel)) and kwargs.get('as_savedmodel')
+
+    def _make_savedmodel(self, obj, serving_input_receiver_fn=None, strip_default_attrs=None):
+        # https://www.tensorflow.org/api_docs/python/tf/keras/experimental/export_saved_model
+        import tensorflow as tf
+        export_dir = tempfile.mkdtemp()
+        tf.contrib.saved_model.save_keras_model(obj, export_dir,
+                                                input_signature=[tf.TensorSpec(shape=obj.input.shape,
+                                                                               dtype=obj.input.dtype,
+                                                                               name='input')])
+        return export_dir
+
+    def fit(self, modelname, Xname, Yname=None, pure_python=True, tpu_specs=None, **kwargs):
+        raise ValueError('cannot fit a saved model')

omegaml/backends/tensorflow/tfsavedmodel.py

@@ -50,12 +50,18 @@ def _extract_savedmodel(self, packagefname):
         rmtree(lpath)
         return model
 
-    def put_model(self, obj, name, attributes=None, serving_input_receiver_fn=None, strip_default_attrs=None):
+    def _make_savedmodel(self, obj, serving_input_receiver_fn=None, strip_default_attrs=None):
         # adapted from https://www.tensorflow.org/guide/saved_model#perform_the_export
         export_dir_base = tempfile.mkdtemp()
         obj.export_savedmodel(export_dir_base,
                               serving_input_receiver_fn=serving_input_receiver_fn,
                               strip_default_attrs=strip_default_attrs)
+        return export_dir_base
+
+    def put_model(self, obj, name, attributes=None, serving_input_receiver_fn=None,
+                  strip_default_attrs=None, **kwargs):
+        export_dir_base = self._make_savedmodel(obj, serving_input_receiver_fn=serving_input_receiver_fn,
+                                                strip_default_attrs=strip_default_attrs)
         zipfname = self._package_savedmodel(export_dir_base, name)
         with open(zipfname, 'rb') as fzip:
             fileid = self.model_store.fs.put(
@@ -107,3 +113,6 @@ def predict(
         if rName:
             result = self.data_store.put(result, rName)
         return result
+
+    def fit(self, modelname, Xname, Yname=None, pure_python=True, tpu_specs=None, **kwargs):
+        raise ValueError('cannot fit a saved model')

omegaml/defaults.py

@@ -47,6 +47,7 @@
     'keras.h5': 'omegaml.backends.keras.KerasBackend',
     'ndarray.bin': 'omegaml.backends.npndarray.NumpyNDArrayBackend',
     'tfkeras.h5': 'omegaml.backends.tensorflow.TensorflowKerasBackend',
+    'tfkeras.savedmodel': 'omegaml.backends.tensorflow.TensorflowKerasSavedModelBackend',
     'tf.savedmodel': 'omegaml.backends.tensorflow.TensorflowSavedModelBackend',
     'tfestimator.model': 'omegaml.backends.tensorflow.TFEstimatorModelBackend',
     'virtualobj.dill': 'omegaml.backends.virtualobj.VirtualObjectBackend',

omegaml/tests/test_tfkerasmodel.py

@@ -3,7 +3,9 @@
 import numpy as np
 
 from omegaml import Omega
+from omegaml.backends.tensorflow import TensorflowSavedModelPredictor
 from omegaml.backends.tensorflow.tfkeras import TensorflowKerasBackend
+from omegaml.backends.tensorflow.tfkerassavedmodel import TensorflowKerasSavedModelBackend
 from omegaml.documents import Metadata
 from omegaml.tests.util import OmegaTestMixin, tf_perhaps_eager_execution
 
@@ -12,6 +14,7 @@ class TensorflowKerasBackendTests(OmegaTestMixin, TestCase):
     def setUp(self):
         self.om = Omega()
         self.om.models.register_backend(TensorflowKerasBackend.KIND, TensorflowKerasBackend)
+        self.om.models.register_backend(TensorflowKerasSavedModelBackend.KIND, TensorflowKerasSavedModelBackend)
         self.clean()
         tf_perhaps_eager_execution()
 
@@ -39,7 +42,7 @@ def _build_model(self, fit=False):
         model.add(Dropout(0.5))
         model.add(Dense(64, activation='relu'))
         model.add(Dropout(0.5))
-        model.add(Dense(10, activation='softmax'))
+        model.add(Dense(10, activation='softmax', name='output'))
         sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
         model.compile(loss='categorical_crossentropy',
                       optimizer=sgd,
@@ -90,6 +93,17 @@ def test_runtime_predict_from_trained_model(self):
         result = om.runtime.model('keras-model').predict(x_test).get()
         self.assertEqual(result.shape, (100, 10))
 
+    def test_save_load_savedmodel(self):
+        om = self.om
+        model = self._build_model(fit=True)
+        x_test = np.random.random((100, 20))
+        yhat = model.predict(x_test)
+        om.models.put(model, 'keras-savedmodel', as_savedmodel=True)
+        model_ = om.models.get('keras-savedmodel')
+        self.assertIsInstance(model_, TensorflowSavedModelPredictor)
+        yhat_ = model_.predict({'input_wrapper_for_input:0': x_test})
+        self.assertTrue(np.allclose(yhat_[model.output_names[0]], yhat))
+
 
 
 

omegaml/tests/test_tfsavedmodel.py

@@ -37,7 +37,7 @@ def _build_model(self):
         # Dense(64) is a fully-connected layer with 64 hidden units.
         # in the first layer, you must specify the expected input data shape:
         # here, 20-dimensional vectors.
-        model.add(Dense(64, activation='relu', input_dim=20, name='features'))
+        model.add(Dense(64, activation='relu', input_dim=20, name='X'))
         model.add(Dropout(0.5))
         model.add(Dense(64, activation='relu'))
         model.add(Dropout(0.5))
@@ -50,16 +50,15 @@ def _build_model(self):
         # https://www.tensorflow.org/guide/estimators
         est_model = tf.keras.estimator.model_to_estimator(keras_model=model)
         train_input_fn = tf.estimator.inputs.numpy_input_fn(
-            x={"features_input": x_train},
+            x={"X_input": x_train},
             y=y_train,
             num_epochs=1,
             shuffle=False)
 
         est_model.train(train_input_fn)
         return est_model
 
-    skipIf(tf_in_eager_execution(), "cannot run in eager mode")
-
+    @skipIf(tf_in_eager_execution(), "cannot run in eager mode")
     def test_save_load(self):
         import tensorflow as tf
         import numpy as np
@@ -69,14 +68,14 @@ def test_save_load(self):
 
         # https://www.tensorflow.org/guide/saved_model#prepare_serving_inputs
         default_batch_size = 1
-        feature_spec = {'features_input': tf.FixedLenFeature(dtype=np.int64, shape=(1,))}
+        feature_spec = {'X_input': tf.FixedLenFeature(dtype=np.float32, shape=(20,))}
 
         def serving_input_receiver_fn():
             """An input receiver that expects a serialized tf.Example."""
             serialized_tf_example = tf.placeholder(dtype=tf.string,
                                                    shape=[default_batch_size],
-                                                   name='input_example_tensor')
-            receiver_tensors = {'examples': serialized_tf_example}
+                                                   name='X_input')
+            receiver_tensors = {'X_input': serialized_tf_example}
             features = tf.parse_example(serialized_tf_example, feature_spec)
             return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
 
@@ -85,3 +84,8 @@ def serving_input_receiver_fn():
         self.assertIn('estimator-savedmodel', om.models.list())
         model_ = om.models.get('estimator-savedmodel')
         self.assertIsInstance(model_, TensorflowSavedModelPredictor)
+        x_test = np.random.random((100, 20))
+        example = tf.train.Example(features=tf.train.Features(feature={
+            'X_input': tf.train.Feature(float_list=tf.train.FloatList(value=x_test[0, :]))
+        }))
+        model_.predict({'X_input': [example.SerializeToString()]})