Repository: sql-machine-learning/elasticdl

Issue:
[Serving] Export the model with elasticdl.layers.Embedding to SavedModel
1. Restore the Keras model def and params by trained params of the elasticdl.layers.Embedding layer.
2. Add inputs and outputs to the model and save the model using SavedModel.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/model_handler.py b/elasticdl/python/common/model_handler.py
--- a/elasticdl/python/common/model_handler.py
+++ b/elasticdl/python/common/model_handler.py
@@ -21,7 +21,7 @@ def get_model_to_train(self, model):
         pass
 
     @abc.abstractmethod
-    def get_model_to_export(self, model, dataset):
+    def get_model_to_export(self, model, trained_params, dataset):
         """
         Get the model which can be exported a SavedModel
         by tf.saved_model.save.
@@ -48,9 +48,14 @@ class DefaultModelHandler(ModelHandler):
     def get_model_to_train(self, model):
         return model
 
-    def get_model_to_export(self, model, dataset):
+    def get_model_to_export(self, model, trained_params, dataset):
+        """
+        Get model with inputs and trained parameters to export.
+        """
         if not model.inputs:
             model._build_model_with_inputs(inputs=dataset, targets=None)
+        for var in model.trainable_variables:
+            var.assign(trained_params[var.name])
         return model
 
 
@@ -61,57 +66,94 @@ def get_model_to_train(self, model):
         an elasticdl.layers.Embedding layer in ParameterServerStrategy.
         """
         if type(model) == tf.keras.Sequential or model._is_graph_network:
-            model = self._clone_model_for_sequential_and_functional(model)
+            model = self._replace_embedding_layer_to_clone_model(
+                model, tf.keras.layers.Embedding, Embedding
+            )
         else:
-            model = self._replace_embedding_attribute_for_subclass(model)
+            model = self._replace_embedding_attribute_for_subclass(
+                model, tf.keras.layers.Embedding, Embedding
+            )
         return model
 
-    def get_model_to_export(self, model, dataset):
+    def get_model_to_export(self, model, trained_params, dataset):
         """
-        To export model for tf-serving by tf.saved_model.save:
-        1. Add inputs and outputs to the model.
-        2. Restore Keras model and replace embedding parameters with trained
-        model.
+        Get the model which can be exported to a SavedModel by
+        `tf.saved_model.save`.
         """
-        # TODO
-        pass
+        model = self._restore_keras_model_def(model)
+        if not model.inputs:
+            # build model to add inputs and outputs that
+            # can be consumed by tf-serving
+            model._build_model_with_inputs(inputs=dataset, targets=None)
+        for var in model.trainable_variables:
+            var.assign(trained_params[var.name])
+        return model
 
-    def _clone_model_for_sequential_and_functional(self, model):
+    def _restore_keras_model_def(self, model):
         """
-        Clone a new model with elasticdl.layers.Embedding for
-        Sequential and functional API model.
+        Restore Keras model definition by replacing
+        `elasticdl.layers.Embedding` layers with
+        `tf.keras.layers.Embedding` layers.
+        """
+        # clear keras model session to avoid clutter from old models/layers.
+        tf.keras.backend.clear_session()
+        if (
+            type(model) == tf.keras.models.Model
+            and not model._is_graph_network
+        ):
+            model = self._replace_embedding_attributes_for_subclass(
+                model, Embedding, tf.keras.layers.Embedding
+            )
+        else:
+            model = self._replace_embedding_layer_to_clone_model(
+                model, Embedding, tf.keras.layers.Embedding
+            )
+        return model
+
+    def _replace_embedding_layer_to_clone_model(
+        self, model, src_embedding_class, dst_embedding_class
+    ):
+        """
+        Clone a new model by cloning model and replace the
+        src_embedding_class layer with a dst_embedding_class.
         """
 
         def _clone_function(layer):
-            if type(layer) == tf.keras.layers.Embedding:
-                logger.info("Replace Keras Embedding with ElasticDL Embedding")
-                edl_embedding_layer = Embedding(
+            if type(layer) == src_embedding_class:
+                logger.debug(
+                    "Replace {} with {}".format(
+                        src_embedding_class, dst_embedding_class
+                    )
+                )
+                embedding_layer = dst_embedding_class(
                     output_dim=layer.output_dim,
                     input_dim=layer.input_dim,
                     embeddings_initializer=layer.embeddings_initializer,
                     mask_zero=layer.mask_zero,
                     input_length=layer.input_length,
                 )
-                return edl_embedding_layer
+                return embedding_layer
             return layer
 
         return tf.keras.models.clone_model(
             model, clone_function=_clone_function
         )
 
-    def _replace_embedding_attributes_for_subclass(self, model):
+    def _replace_embedding_attributes_for_subclass(
+        self, model, src_embedding_class, dst_embedding_class
+    ):
         """
         Replace the keras embedding attribute with
         elasticdl.layers.Embedding layer.
         """
-        for attr_name, attr_value in model.__dict__.items():
-            if type(attr_value) == tf.keras.layers.Embedding:
-                edl_embedding_layer = Embedding(
-                    output_dim=attr_value.output_dim,
-                    input_dim=attr_value.input_dim,
-                    embedding_initializer=attr_value.embeddings_initializer,
-                    mask_zero=attr_value.mask_zero,
-                    input_length=attr_value.input_length,
+        for name, value in model.__dict__.items():
+            if type(value) == src_embedding_class:
+                embedding_layer = dst_embedding_class(
+                    output_dim=value.output_dim,
+                    input_dim=value.input_dim,
+                    embeddings_initializer=value.embeddings_initializer,
+                    mask_zero=value.mask_zero,
+                    input_length=value.input_length,
                 )
-                setattr(model, attr_name, edl_embedding_layer)
+                setattr(model, name, embedding_layer)
         return model
diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py
--- a/elasticdl/python/elasticdl/layers/embedding.py
+++ b/elasticdl/python/elasticdl/layers/embedding.py
@@ -55,6 +55,7 @@ def __init__(
         self.input_dim = input_dim
         self.output_dim = output_dim
         self.embeddings_initializer = embeddings_initializer
+        self.mask_zero = mask_zero
         self.supports_masking = mask_zero
         self.input_length = input_length
         self.combiner = combiner


Test patch:
diff --git a/elasticdl/python/tests/model_handler_test.py b/elasticdl/python/tests/model_handler_test.py
--- a/elasticdl/python/tests/model_handler_test.py
+++ b/elasticdl/python/tests/model_handler_test.py
@@ -32,6 +32,13 @@ def _get_dataset():
     return dataset
 
 
+def _mock_model_trained_params(model):
+    trained_params = {}
+    for var in model.trainable_variables:
+        trained_params[var.name] = np.ones(var.shape.as_list())
+    return trained_params
+
+
 class DefaultModelHandlerTest(unittest.TestCase):
     def setUp(self):
         self.model_handler = ModelHandler.get_model_handler()
@@ -42,16 +49,34 @@ def test_get_model_to_ps(self):
         self.assertEqual(type(model_inst.layers[1]), tf.keras.layers.Embedding)
 
     def test_get_model_to_export(self):
+        dataset = _get_dataset()
         feature_columns = feature_columns_fn()
         model_inst = custom_sequential_model(feature_columns)
-        dataset = _get_dataset()
-        self.model_handler.get_model_to_export(model_inst, dataset)
+        model_inst._build_model_with_inputs(inputs=dataset, targets=None)
+        trained_params = _mock_model_trained_params(model_inst)
+        model_inst = self.model_handler.get_model_to_export(
+            model_inst, trained_params, dataset
+        )
         self.assertEqual(list(model_inst.inputs.keys()), ["age", "education"])
         self.assertEqual(len(model_inst.outputs), 1)
 
+        test_data = {
+            "age": [14, 56, 78, 38, 80],
+            "education": [
+                "Bachelors",
+                "Master",
+                "Some-college",
+                "Bachelors",
+                "Master",
+            ],
+        }
+        result = model_inst.call(test_data).numpy()
+        self.assertEqual(result.tolist(), np.ones((5, 1)).tolist())
+
 
 class ParameterSeverModelHandlerTest(unittest.TestCase):
     def setUp(self):
+        tf.keras.backend.clear_session()
         self.model_handler = ModelHandler.get_model_handler(
             distribution_strategy="ParameterServerStrategy"
         )
@@ -62,7 +87,17 @@ def test_get_model_to_train(self):
         self.assertEqual(type(model_inst.layers[1]), Embedding)
 
     def test_get_model_to_export(self):
-        pass
+        model_inst = custom_model_with_embedding()
+        trained_params = _mock_model_trained_params(model_inst)
+
+        train_model = self.model_handler.get_model_to_train(model_inst)
+        export_model = self.model_handler.get_model_to_export(
+            train_model, trained_params, dataset=None
+        )
+
+        test_data = tf.constant([0])
+        result = export_model.call(test_data).numpy()
+        self.assertEqual(result[0][0], 3.0)
 
 
 if __name__ == "__main__":