Repository: sql-machine-learning/elasticdl

Issue:
add subclass test case for model handle test 
We only define a functional model in test_module.
https://github.com/sql-machine-learning/elasticdl/blob/35d0647a737d93e6969764cabe541ac8078d13bf/elasticdl/python/tests/test_module.py#L16-L20

We should add subclass one to ensure ModelHandle logic correct.


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
@@ -128,7 +128,7 @@ def _restore_keras_model_def(self, model):
         # clear keras model session to avoid clutter from old models/layers.
         tf.keras.backend.clear_session()
         if (
-            type(model) == tf.keras.models.Model
+            isinstance(model, tf.keras.models.Model)
             and not model._is_graph_network
         ):
             model = self._replace_embedding_attributes_for_subclass(


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
@@ -1,4 +1,3 @@
-import os
 import unittest
 
 import numpy as np
@@ -8,13 +7,45 @@
 from elasticdl.python.elasticdl.layers.embedding import Embedding
 from elasticdl.python.master.checkpoint_service import CheckpointService
 from elasticdl.python.master.servicer import MasterServicer
-from elasticdl.python.tests.test_module import (
-    custom_model_with_embedding,
-    custom_sequential_model,
-    feature_columns_fn,
-)
 
-_model_zoo_path = os.path.dirname(os.path.realpath(__file__))
+
+class CustomModel(tf.keras.models.Model):
+    def __init__(self):
+        super(CustomModel, self).__init__()
+        self.embedding = tf.keras.layers.Embedding(4, 2)
+        self.dense = tf.keras.layers.Dense(1)
+
+    def call(self, inputs):
+        embedding = self.embedding(inputs)
+        output = self.dense(embedding)
+        return output
+
+
+def custom_model_with_embedding():
+    inputs = tf.keras.layers.Input(shape=(4,), name="x")
+    embedding = tf.keras.layers.Embedding(4, 2)(inputs)
+    outputs = tf.keras.layers.Dense(1)(embedding)
+    return tf.keras.models.Model(inputs, outputs)
+
+
+def custom_sequential_model(feature_columns):
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.DenseFeatures(feature_columns=feature_columns),
+            tf.keras.layers.Dense(10, activation="relu"),
+            tf.keras.layers.Dense(1, activation="sigmoid"),
+        ]
+    )
+    return model
+
+
+def feature_columns_fn():
+    age = tf.feature_column.numeric_column("age", dtype=tf.int64)
+    education = tf.feature_column.categorical_column_with_hash_bucket(
+        "education", hash_bucket_size=4
+    )
+    education_one_hot = tf.feature_column.indicator_column(education)
+    return [age, education_one_hot]
 
 
 def _get_dataset():
@@ -119,6 +150,39 @@ def test_get_model_to_export(self):
         result = export_model.call(test_data).numpy()
         self.assertEqual(result[0][0], 3.0)
 
+    def test_get_subclass_model_to_export(self):
+        def _get_dataset():
+            dataset = tf.data.Dataset.from_tensor_slices(
+                np.random.randint(0, 10, (10, 4))
+            )
+            dataset = dataset.batch(2)
+            return dataset
+
+        model_inst = CustomModel()
+        dataset = _get_dataset()
+
+        trained_params = {
+            "custom_model/embedding/embeddings:0": np.ones(
+                (4, 2), dtype="float32"
+            ),
+            "custom_model/dense/kernel:0": np.ones((2, 1), dtype="float32"),
+            "custom_model/dense/bias:0": np.ones((1), dtype="float32"),
+        }
+
+        for name, value in trained_params.items():
+            self.master.set_model_var(name, value)
+
+        train_model = self.model_handler.get_model_to_train(model_inst)
+        self.assertEqual(type(train_model.embedding), Embedding)
+
+        export_model = self.model_handler.get_model_to_export(
+            train_model, dataset=dataset
+        )
+
+        test_data = tf.constant([0])
+        result = export_model.call(test_data).numpy()
+        self.assertEqual(result[0][0], 3.0)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/elasticdl/python/tests/test_module.py b/elasticdl/python/tests/test_module.py
--- a/elasticdl/python/tests/test_module.py
+++ b/elasticdl/python/tests/test_module.py
@@ -1,6 +1,6 @@
 import tensorflow as tf
 from tensorflow.keras import Model
-from tensorflow.keras.layers import Dense, Embedding, Input
+from tensorflow.keras.layers import Dense, Input
 
 from elasticdl.python.worker.prediction_outputs_processor import (
     BasePredictionOutputsProcessor,
@@ -13,33 +13,6 @@ def custom_model():
     return Model(inputs, outputs)
 
 
-def custom_model_with_embedding():
-    inputs = Input(shape=(4,), name="x")
-    embedding = Embedding(4, 2)(inputs)
-    outputs = Dense(1)(embedding)
-    return Model(inputs, outputs)
-
-
-def custom_sequential_model(feature_columns):
-    model = tf.keras.Sequential(
-        [
-            tf.keras.layers.DenseFeatures(feature_columns=feature_columns),
-            tf.keras.layers.Dense(10, activation="relu"),
-            tf.keras.layers.Dense(1, activation="sigmoid"),
-        ]
-    )
-    return model
-
-
-def feature_columns_fn():
-    age = tf.feature_column.numeric_column("age", dtype=tf.int64)
-    education = tf.feature_column.categorical_column_with_hash_bucket(
-        "education", hash_bucket_size=4
-    )
-    education_one_hot = tf.feature_column.indicator_column(education)
-    return [age, education_one_hot]
-
-
 def loss(predictions, labels):
     return tf.reduce_mean(tf.square(predictions - labels))