Repository: sql-machine-learning/elasticdl

Issue:
[Serving] Export the model with ElasticDL embedding_column to the TF native SavedModel.



Discussion:
The demo to replace the embedding_column weights is:
```python
import time
import numpy as np
import tensorflow as tf 

def mock_embedding_column_weights(feature_columns, dense_feature_name):
    embedding_column_weights = {}
    for fc in feature_columns:
        if isinstance(fc, EmbeddingColumn):
            variable_name = "/".join([dense_feature_name, fc.name, "embedding_weights:0"])
            weight_shape = (fc.categorical_column.num_buckets, fc.dimension)
            embedding_column_weights[variable_name] = np.ones(weight_shape)
    return embedding_column_weights


def get_feature_columns():
    age = tf.feature_column.numeric_column("age", dtype=tf.int64)
    edu_embedding = tf.feature_column.embedding_column(
        tf.feature_column.categorical_column_with_hash_bucket(
        'education', hash_bucket_size=4),
        4
    )
    edu_embedding_1 = tf.feature_column.embedding_column(
        tf.feature_column.categorical_column_with_hash_bucket(
        'education_1', hash_bucket_size=4),
        1
    )
    return [age, edu_embedding, edu_embedding_1]


def custom_model(feature_columns):
    input_layers = {}
    input_layers['age'] = tf.keras.layers.Input(name='age', shape=(1,))
    input_layers['education'] = tf.keras.layers.Input(name='education', shape=(1,), dtype=tf.string)
    
    input_layers['education_1'] = tf.keras.layers.Input(name='education_1', shape=(1,), dtype=tf.string)
    
    dense_feature = tf.keras.layers.DenseFeatures(feature_columns=feat_cols)(input_layers)
    return tf.keras.models.Model(inputs=input_layers, outputs=dense_feature)

feat_cols = get_feature_columns()
model = custom_model(feat_cols)
output = model.call(
    {
        'age':tf.constant([[10],[16]]),
        'education':tf.constant([['Bachelors'],['Master']]),
        'education_1':tf.constant([['Bachelors'],['Master']])
    }
)

print("Before assign: ", output)


model_feature_columns = None
for layer in model.layers:
    if isinstance(layer, tf.keras.layers.DenseFeatures):
        model_feature_columns = list(layer._feature_columns)  
        embedding_column_weights = mock_embedding_column_weights(
            model_feature_columns, layer.name
        )
for weight in model.trainable_weights:
    if weight.name in embedding_column_weights:
        weight.assign(embedding_column_weights[weight.name])
output = model.call(
    {
        'age':tf.constant([[10],[16]]),
        'education':tf.constant([['Bachelors'],['Master']]),
        'education_1':tf.constant([['Bachelors'],['Master']]),
    }
)

print("After assign: ", output)
```

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
@@ -8,6 +8,7 @@
 from elasticdl.python.common.log_utils import default_logger as logger
 from elasticdl.python.common.save_utils import CheckpointSaver
 from elasticdl.python.elasticdl.feature_column.feature_column import (
+    EmbeddingColumn,
     embedding_column,
 )
 from elasticdl.python.elasticdl.layers.embedding import Embedding
@@ -111,6 +112,26 @@ def _replace_tf_embedding_column_with_edl(dense_features_layer):
     )
 
 
+def _replace_edl_embedding_column_with_tf(dense_features_layer):
+    new_feature_columns = []
+    for column in dense_features_layer._feature_columns:
+        if isinstance(column, EmbeddingColumn):
+            logger.info(
+                "Replace embedding_column {} from ElasticDL "
+                "version to TF version".format(column.name)
+            )
+            new_column = fc_lib.embedding_column(
+                column.categorical_column, dimension=column.dimension
+            )
+            new_feature_columns.append(new_column)
+        else:
+            new_feature_columns.append(column)
+
+    return tf.keras.layers.DenseFeatures(
+        feature_columns=new_feature_columns, name=dense_features_layer.name
+    )
+
+
 class ModelHandler(metaclass=abc.ABCMeta):
     """Generate the model to train in ElasticDL for different distributed
     strategies and export trained model in ElasticDL to SavedModel.
@@ -334,6 +355,8 @@ def _clone_function(layer):
                         name=layer.name,
                     )
                 return embedding_layer
+            elif type(layer) == tf.keras.layers.DenseFeatures:
+                return _replace_edl_embedding_column_with_tf(layer)
             return layer
 
         return tf.keras.models.clone_model(
@@ -424,4 +447,7 @@ def _replace_attr_with_keras_embedding(model):
                         input_length=value.input_length,
                     )
                 setattr(model, name, embedding_layer)
+            elif type(value) == tf.keras.layers.DenseFeatures:
+                feature_layer = _replace_edl_embedding_column_with_tf(value)
+                setattr(model, name, feature_layer)
         return model


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
@@ -14,6 +14,7 @@
 )
 from elasticdl.python.elasticdl.layers.embedding import Embedding
 from elasticdl.python.keras.layers import SparseEmbedding
+from elasticdl.python.ps.embedding_table import EmbeddingTable
 from elasticdl.python.ps.parameters import Parameters
 
 EMBEDDING_INPUT_DIM = 300000
@@ -39,14 +40,25 @@ def custom_model_with_embedding_layer():
 
 
 def custom_model_with_embedding_column():
+    inputs = {
+        "age": tf.keras.layers.Input(shape=(1,), name="age", dtype=tf.float32),
+        "user_id": tf.keras.layers.Input(
+            shape=(1,), name="user_id", dtype=tf.string
+        ),
+    }
     age = tf.feature_column.numeric_column("age", dtype=tf.int64)
     user_id_embedding = tf.feature_column.embedding_column(
         tf.feature_column.categorical_column_with_hash_bucket(
             "user_id", hash_bucket_size=EMBEDDING_INPUT_DIM
         ),
-        dimension=8,
+        dimension=2,
+    )
+    feature_columns = [age, user_id_embedding]
+    dense = tf.keras.layers.DenseFeatures(feature_columns=feature_columns)(
+        inputs
     )
-    return custom_sequential_model([age, user_id_embedding])
+    output = tf.keras.layers.Dense(1)(dense)
+    return tf.keras.models.Model(inputs, output)
 
 
 def custom_model_with_sparse_embedding():
@@ -152,21 +164,30 @@ def setUp(self):
             checkpoint_dir="",
         )
 
-    def _save_model(self, is_subclass):
-        prefix = "custom_model/" if is_subclass else ""
-        ckpt_dir = self.model_handler._checkpoint_dir
-        checkpoint_saver = CheckpointSaver(ckpt_dir, 0, 0, False)
+    def _mock_model_parameters(self, model):
         params = Parameters()
-        params.non_embedding_params[
-            prefix + "embedding/embeddings:0"
-        ] = tf.Variable(tf.ones([EMBEDDING_INPUT_DIM, 2]), dtype=tf.float32)
-        params.non_embedding_params[prefix + "dense/kernel:0"] = tf.Variable(
-            [[1.0], [1.0]]
-        )
-        params.non_embedding_params[prefix + "dense/bias:0"] = tf.Variable(
-            [1.0]
-        )
+        for weight in model.trainable_variables:
+            if "embedding" in weight.name:
+                embedding_table = EmbeddingTable(
+                    name=weight.name,
+                    dim=weight.shape[1],
+                    initializer="RandomUniform",
+                )
+                embedding_table.set(
+                    np.arange(weight.shape[0]), np.ones(weight.shape)
+                )
+                params.embedding_params[weight.name] = embedding_table
+            else:
+                params.non_embedding_params[weight.name] = tf.ones(
+                    weight.shape
+                )
         params.version = 100
+        return params
+
+    def _mock_model_weights_and_save_checkpoint(self, model):
+        ckpt_dir = self.model_handler._checkpoint_dir
+        checkpoint_saver = CheckpointSaver(ckpt_dir, 0, 0, False)
+        params = self._mock_model_parameters(model)
         model_pb = params.to_model_pb()
         checkpoint_saver.save(100, model_pb, False)
 
@@ -177,23 +198,52 @@ def test_get_model_with_embedding_layer_to_train(self):
 
     def test_get_model_with_embedding_column_to_train(self):
         model_inst = custom_model_with_embedding_column()
+        dense_features_layer_index = 2
+        embedding_column_index = 1
         self.assertEqual(
-            type(model_inst.layers[0]._feature_columns[1]),
+            type(
+                model_inst.layers[dense_features_layer_index]._feature_columns[
+                    embedding_column_index
+                ]
+            ),
             fc_lib.EmbeddingColumn,
         )
         model_inst = self.model_handler.get_model_to_train(model_inst)
         self.assertEqual(
-            type(model_inst.layers[0]._feature_columns[1]), EmbeddingColumn
+            type(
+                model_inst.layers[dense_features_layer_index]._feature_columns[
+                    embedding_column_index
+                ]
+            ),
+            EmbeddingColumn,
         )
 
+    def test_get_model_with_embedding_column_to_export(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            self.model_handler._checkpoint_dir = os.path.join(
+                temp_dir, "test_export"
+            )
+            model_inst = custom_model_with_embedding_column()
+            self._mock_model_weights_and_save_checkpoint(model_inst)
+
+            model_inst = self.model_handler.get_model_to_train(model_inst)
+            export_model = self.model_handler.get_model_to_export(
+                model_inst, dataset=None
+            )
+            result = export_model.call(
+                {"age": tf.constant([[1]]), "user_id": tf.constant([["134"]])}
+            ).numpy()
+            self.assertEqual(result[0][0], 4.0)
+
     def test_get_model_to_export(self):
         with tempfile.TemporaryDirectory() as temp_dir:
             self.model_handler._checkpoint_dir = os.path.join(
                 temp_dir, "test_export"
             )
-            self._save_model(False)
             model_inst = custom_model_with_embedding_layer()
             train_model = self.model_handler.get_model_to_train(model_inst)
+
+            self._mock_model_weights_and_save_checkpoint(model_inst)
             export_model = self.model_handler.get_model_to_export(
                 train_model, dataset=None
             )
@@ -207,7 +257,6 @@ def test_get_subclass_model_to_export(self):
             self.model_handler._checkpoint_dir = os.path.join(
                 temp_dir, "test_export"
             )
-            self._save_model(True)
 
             def _get_dataset():
                 dataset = tf.data.Dataset.from_tensor_slices(
@@ -218,7 +267,10 @@ def _get_dataset():
 
             model_inst = CustomModel()
             dataset = _get_dataset()
+            model_inst._build_model_with_inputs(inputs=dataset, targets=None)
+            self._mock_model_weights_and_save_checkpoint(model_inst)
 
+            model_inst.inputs = None  # Reset model inputs
             train_model = self.model_handler.get_model_to_train(model_inst)
             self.assertEqual(type(train_model.embedding), Embedding)
 
@@ -240,10 +292,10 @@ def test_get_model_with_sparse_to_export(self):
             self.model_handler._checkpoint_dir = os.path.join(
                 temp_dir, "test_export"
             )
-            self._save_model(False)
             model_inst = custom_model_with_sparse_embedding()
             train_model = self.model_handler.get_model_to_train(model_inst)
 
+            self._mock_model_weights_and_save_checkpoint(model_inst)
             # Model handler will restore model parameters from the checkpoint
             # directory and assign parameters to train_model.
             export_model = self.model_handler.get_model_to_export(