Repository: sql-machine-learning/elasticdl

Issue:
DeepFM unit test for new PS



Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/tensor.py b/elasticdl/python/common/tensor.py
--- a/elasticdl/python/common/tensor.py
+++ b/elasticdl/python/common/tensor.py
@@ -6,7 +6,6 @@
     dtype_numpy_to_tensor,
     dtype_tensor_to_numpy,
 )
-from elasticdl.python.common.log_utils import default_logger as logger
 
 
 class Tensor(object):
@@ -53,13 +52,8 @@ def set(self, values=None, indices=None, name=None):
                     "tf.IndexedSlices, indices must be None."
                 )
             if values.dense_shape is not None:
-                # TODO(yunjian.lmh): Support dense shape, or do not print
-                #     warning message, or there will be too much warning
-                #     messages.
-                logger.warning(
-                    "ElasticDL Tensor ignores dense_shape in "
-                    "TensorFlow.IndexedSlices."
-                )
+                # TODO(yunjian.lmh): Support dense shape
+                pass
 
             self.values = values.values.numpy()
             self.indices = values.indices.numpy()


Test patch:
diff --git a/elasticdl/python/tests/pserver_servicer_test.py b/elasticdl/python/tests/pserver_servicer_test.py
--- a/elasticdl/python/tests/pserver_servicer_test.py
+++ b/elasticdl/python/tests/pserver_servicer_test.py
@@ -61,7 +61,7 @@ def create_server_and_stub(
             port=self._port,
             model_zoo=_test_model_zoo_path,
             model_def="test_module.custom_model",
-            **kwargs,
+            **kwargs
         )
         pserver = ParameterServer(args)
         pserver.prepare()
@@ -259,18 +259,18 @@ def push_gradient_test_setup(self):
         self.embedding_table = (
             np.random.rand(4 * dim).reshape((4, dim)).astype(np.float32)
         )
-        self.embedding_grads = tf.IndexedSlices(
+        self.embedding_grads0 = tf.IndexedSlices(
             values=np.random.rand(3 * dim)
             .reshape((3, dim))
             .astype(np.float32),
             indices=(3, 1, 3),
         )
-        self.expected_embed_table = np.copy(self.embedding_table)
-        for gv, gi in zip(
-            self.embedding_grads.values, self.embedding_grads.indices
-        ):
-            self.expected_embed_table[gi] -= self._lr * gv
-
+        self.embedding_grads1 = tf.IndexedSlices(
+            values=np.random.rand(3 * dim)
+            .reshape((3, dim))
+            .astype(np.float32),
+            indices=(2, 2, 3),
+        )
         push_model_req = elasticdl_pb2.Model()
         push_model_req.version = self._parameters.version
         for name, value in zip(self.var_names, self.var_values):
@@ -297,8 +297,8 @@ def test_push_gradient_async_update(self):
             emplace_tensor_pb_from_ndarray(req.gradients, g, name=name)
         emplace_tensor_pb_from_ndarray(
             req.gradients,
-            values=self.embedding_grads.values,
-            indices=self.embedding_grads.indices,
+            values=self.embedding_grads0.values,
+            indices=self.embedding_grads0.indices,
             name=self._embedding_info.name,
         )
         res = self._stub.push_gradient(req)
@@ -316,12 +316,16 @@ def test_push_gradient_async_update(self):
                 )
             )
 
+        expected_embed_table = np.copy(self.embedding_table)
+        for gv, gi in zip(
+            self.embedding_grads0.values, self.embedding_grads0.indices
+        ):
+            expected_embed_table[gi] -= self._lr * gv
+
         actual_embed_table = self._parameters.get_embedding_param(
-            self._embedding_info.name, range(len(self.expected_embed_table))
-        )
-        self.assertTrue(
-            np.allclose(self.expected_embed_table, actual_embed_table)
+            self._embedding_info.name, range(len(expected_embed_table))
         )
+        self.assertTrue(np.allclose(expected_embed_table, actual_embed_table))
 
         # Test applying gradients with same name
         for name, var in zip(self.var_names, self.var_values):
@@ -358,6 +362,12 @@ def test_push_gradient_sync_update(self):
         req.model_version = 0
         for g, name in zip(self.grad_values0, self.var_names):
             emplace_tensor_pb_from_ndarray(req.gradients, g, name=name)
+        emplace_tensor_pb_from_ndarray(
+            req.gradients,
+            values=self.embedding_grads0.values,
+            indices=self.embedding_grads0.indices,
+            name=self._embedding_info.name,
+        )
         res = self._stub.push_gradient(req)
         self.assertEqual(res.accepted, True)
         self.assertEqual(res.model_version, 0)
@@ -368,8 +378,8 @@ def test_push_gradient_sync_update(self):
             emplace_tensor_pb_from_ndarray(req.gradients, g, name=name)
         emplace_tensor_pb_from_ndarray(
             req.gradients,
-            values=self.embedding_grads.values,
-            indices=self.embedding_grads.indices,
+            values=self.embedding_grads1.values,
+            indices=self.embedding_grads1.indices,
             name=self._embedding_info.name,
         )
         res = self._stub.push_gradient(req)
@@ -398,9 +408,21 @@ def test_push_gradient_sync_update(self):
                 )
             )
 
+        expected_embed_table = np.copy(self.embedding_table)
+        for gv, gi in zip(
+            self.embedding_grads0.values, self.embedding_grads0.indices
+        ):
+            expected_embed_table[gi] -= self._lr * gv
+        for gv, gi in zip(
+            self.embedding_grads1.values, self.embedding_grads1.indices
+        ):
+            expected_embed_table[gi] -= self._lr * gv
+
         actual_embed_table = self._parameters.get_embedding_param(
-            self._embedding_info.name, range(len(self.expected_embed_table))
-        )
-        self.assertTrue(
-            np.allclose(self.expected_embed_table, actual_embed_table)
+            self._embedding_info.name, range(len(expected_embed_table))
         )
+        self.assertTrue(np.allclose(expected_embed_table, actual_embed_table))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/elasticdl/python/tests/worker_ps_interaction_test.py b/elasticdl/python/tests/worker_ps_interaction_test.py
--- a/elasticdl/python/tests/worker_ps_interaction_test.py
+++ b/elasticdl/python/tests/worker_ps_interaction_test.py
@@ -1,5 +1,6 @@
 import os
 import unittest
+from pathlib import Path
 
 import grpc
 import numpy as np
@@ -9,6 +10,9 @@
 from elasticdl.python.common.constants import GRPC
 from elasticdl.python.common.hash_utils import int_to_id, string_to_id
 from elasticdl.python.common.model_utils import get_model_spec
+from elasticdl.python.data.recordio_gen.frappe_recordio_gen import (
+    load_raw_data,
+)
 from elasticdl.python.ps.embedding_table import EmbeddingTable
 from elasticdl.python.ps.parameter_server import ParameterServer
 from elasticdl.python.tests.test_utils import PserverArgs
@@ -74,6 +78,45 @@ def _create_pserver_and_channel(self, ports):
             channels.append(channel)
         return pservers, channels
 
+    def test_worker_pull_embedding(self):
+        worker = Worker(
+            worker_id=0,
+            job_type=elasticdl_pb2.TRAINING,
+            minibatch_size=self._batch_size,
+            model_zoo=self._model_zoo_path,
+            model_def=self._model_def,
+            ps_channels=self._channel,
+        )
+
+        # Test lookup embedding vectors that do not exist
+        layers = ["test-2", "test-2-slot"]
+        ids = [3, 5, 1, 6, 10, 2, 1, 2, 4, 7, 9]
+        embedding_table_args = [
+            (layers[0], 8, "uniform", False),
+            (layers[1], 8, 3.3, True),
+        ]
+
+        # initialize embedding table object
+        for pserver in self._pserver:
+            for layer, table_args in zip(layers, embedding_table_args):
+                pserver.parameters.embedding_params[layer] = EmbeddingTable(
+                    *table_args
+                )
+
+        result_dict = {}
+        for layer in layers:
+            embedding = worker.pull_embedding_vector(layer, ids)
+            result_dict[layer] = embedding
+
+        for layer in layers:
+            expected_result = []
+            for embedding_id in ids:
+                ps_id = int_to_id(embedding_id, len(self._pserver))
+                table = self._pserver[ps_id].parameters.embedding_params[layer]
+                expected_result.append(table.get([embedding_id]))
+            expected_result = np.concatenate(expected_result)
+            self.assertTrue(np.allclose(expected_result, result_dict[layer]))
+
     def test_compare_onebatch_train(self):
         images, labels = random_batch(self._batch_size)
         # TODO(yunjian.lmh): test optimizer wrapper
@@ -84,7 +127,9 @@ def test_compare_onebatch_train(self):
             job_type=elasticdl_pb2.TRAINING,
             minibatch_size=self._batch_size,
             model_zoo=self._model_zoo_path,
-            model_def=self._model_def,
+            model_def=(
+                "mnist_functional_api.mnist_functional_api.custom_model"
+            ),
             ps_channels=self._channel,
         )
         worker._run_model_call_before_training(images)
@@ -104,7 +149,9 @@ def test_compare_onebatch_train(self):
             prediction_outputs_processor,
         ) = get_model_spec(
             model_zoo=self._model_zoo_path,
-            model_def=self._model_def,
+            model_def=(
+                "mnist_functional_api.mnist_functional_api.custom_model"
+            ),
             dataset_fn="dataset_fn",
             model_params=None,
             loss="loss",
@@ -127,38 +174,28 @@ def test_compare_onebatch_train(self):
             )
             np.testing.assert_array_equal(ps_v.numpy(), v.numpy())
 
-    def test_compare_mnist_train(self):
-        (
-            (x_train, y_train),
-            (x_test, y_test),
-        ) = tf.keras.datasets.mnist.load_data()
-        x_train = tf.convert_to_tensor(x_train, dtype=tf.float32) / 255.0
-        y_train = tf.convert_to_tensor(y_train, dtype=tf.int32)
-
-        x_test = tf.convert_to_tensor(x_test, dtype=tf.float32) / 255.0
-        y_test = tf.convert_to_tensor(y_test, dtype=tf.int32)
-
-        db = tf.data.Dataset.from_tensor_slices((x_train, y_train))
-        db = db.batch(self._batch_size).repeat(10)
-        test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
-        test_db = test_db.batch(self._batch_size)
-
-        acc_meter = tf.keras.metrics.Accuracy()
-
-        tf.keras.backend.clear_session()
-        tf.random.set_seed(22)
-
+    def _worker_train(self, train_db, test_db, dataset, stop_step):
+        if dataset == "mnist":
+            model_def = (
+                "mnist_functional_api.mnist_functional_api.custom_model"
+            )
+        elif dataset == "frappe":
+            model_def = (
+                "deepfm_functional_api.deepfm_functional_api.custom_model"
+            )
+        else:
+            raise ValueError("dataset %s is not supported", dataset)
         worker = Worker(
             worker_id=0,
             job_type=elasticdl_pb2.TRAINING,
             minibatch_size=self._batch_size,
             model_zoo=self._model_zoo_path,
-            model_def=self._model_def,
+            model_def=model_def,
             ps_channels=self._channel,
         )
-
+        acc_meter = tf.keras.metrics.Accuracy()
         worker_results = []
-        for step, (x, y) in enumerate(db):
+        for step, (x, y) in enumerate(train_db):
             if step == 0:
                 worker._run_model_call_before_training(x)
                 worker.report_variable()
@@ -171,19 +208,53 @@ def test_compare_mnist_train(self):
                 worker.get_model(step, elasticdl_pb2.MINIMUM)
                 for (x, y) in test_db:
                     out = worker.forward_process(x)
-                    acc_meter.update_state(tf.argmax(out, axis=1), y)
-
+                    if dataset == "mnist":
+                        acc_meter.update_state(tf.argmax(out, axis=1), y)
+                    else:
+                        out["probs"] = tf.reshape(out["probs"], [-1])
+                        acc_meter.update_state(
+                            tf.where(
+                                out["probs"] < 0.5,
+                                x=tf.zeros_like(y),
+                                y=tf.ones_like(y),
+                            ),
+                            y,
+                        )
                 worker_results.append(
                     (float(w_loss.numpy()), float(acc_meter.result().numpy()))
                 )
                 acc_meter.reset_states()
 
-            if step > 50:
+            if step > stop_step:
                 break
+        return worker_results
+
+    def test_compare_mnist_train(self):
+        (
+            (x_train, y_train),
+            (x_test, y_test),
+        ) = tf.keras.datasets.mnist.load_data()
+        x_train = tf.convert_to_tensor(x_train, dtype=tf.float32) / 255.0
+        y_train = tf.convert_to_tensor(y_train, dtype=tf.int32)
+
+        x_test = tf.convert_to_tensor(x_test, dtype=tf.float32) / 255.0
+        y_test = tf.convert_to_tensor(y_test, dtype=tf.int32)
+
+        db = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+        db = db.batch(self._batch_size).repeat(10)
+        test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+        test_db = test_db.batch(self._batch_size)
 
         tf.keras.backend.clear_session()
         tf.random.set_seed(22)
-        acc_meter.reset_states()
+        worker_results = self._worker_train(
+            train_db=db, test_db=test_db, dataset="mnist", stop_step=50
+        )
+
+        tf.keras.backend.clear_session()
+        tf.random.set_seed(22)
+
+        acc_meter = tf.keras.metrics.Accuracy()
 
         (
             model,
@@ -194,7 +265,9 @@ def test_compare_mnist_train(self):
             prediction_outputs_processor,
         ) = get_model_spec(
             model_zoo=self._model_zoo_path,
-            model_def=self._model_def,
+            model_def=(
+                "mnist_functional_api.mnist_functional_api.custom_model"
+            ),
             dataset_fn="dataset_fn",
             model_params=None,
             loss="loss",
@@ -227,44 +300,33 @@ def test_compare_mnist_train(self):
         for w, l in zip(worker_results, local_results):
             self.assertTupleEqual(w, l)
 
-    def test_worker_pull_embedding(self):
-        worker = Worker(
-            worker_id=0,
-            job_type=elasticdl_pb2.TRAINING,
-            minibatch_size=self._batch_size,
-            model_zoo=self._model_zoo_path,
-            model_def=self._model_def,
-            ps_channels=self._channel,
-        )
+    def test_deepfm_train(self):
+        tf.keras.backend.clear_session()
+        tf.random.set_seed(22)
+        home = str(Path.home())
 
-        # Test lookup embedding vectors that do not exist
-        layers = ["test-2", "test-2-slot"]
-        ids = [3, 5, 1, 6, 10, 2, 1, 2, 4, 7, 9]
-        embedding_table_args = [
-            (layers[0], 8, "uniform", False),
-            (layers[1], 8, 3.3, True),
-        ]
+        class TmpArgs(object):
+            def __init__(self, data):
+                self.data = data
 
-        # initialize embedding table object
-        for pserver in self._pserver:
-            for layer, table_args in zip(layers, embedding_table_args):
-                pserver.parameters.embedding_params[layer] = EmbeddingTable(
-                    *table_args
-                )
+        args = TmpArgs(data=home + "/.keras/datasets/")
 
-        result_dict = {}
-        for layer in layers:
-            embedding = worker.pull_embedding_vector(layer, ids)
-            result_dict[layer] = embedding
+        x_train, y_train, x_val, y_val, x_test, y_test = load_raw_data(args)
+        x_train = tf.convert_to_tensor(x_train, dtype=tf.int64)
+        x_test = tf.convert_to_tensor(x_test, dtype=tf.int64)
+        y_train = tf.convert_to_tensor(y_train, dtype=tf.int64)
+        y_test = tf.convert_to_tensor(y_test, dtype=tf.int64)
 
-        for layer in layers:
-            expected_result = []
-            for embedding_id in ids:
-                ps_id = int_to_id(embedding_id, len(self._pserver))
-                table = self._pserver[ps_id].parameters.embedding_params[layer]
-                expected_result.append(table.get([embedding_id]))
-            expected_result = np.concatenate(expected_result)
-            self.assertTrue(np.allclose(expected_result, result_dict[layer]))
+        db = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+        db = db.batch(self._batch_size).repeat(10)
+        test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+        test_db = test_db.batch(self._batch_size)
+
+        worker_results = self._worker_train(
+            train_db=db, test_db=test_db, dataset="frappe", stop_step=200
+        )
+        acc = max([r[1] for r in worker_results])
+        self.assertLess(0.7, acc)
 
 
 if __name__ == "__main__":