Repository: sql-machine-learning/elasticdl

Issue:
Extend `Worker` to support model evaluation task
Part of https://github.com/wangkuiyi/elasticdl/issues/384. 

Extend `Worker` to support the following:
    * `distributed_evaluate()` that contains the main logic for model evaluation.
    * `report_task_result()` that reports evaluation task result (e.g. task id and error message) back to master through RPC call.
    * `report_evaluation_metrics()` that reports the computed evaluation metrics (e.g. accuracy, precision, recall, etc.) back to master through RPC call.
Extend `Worker` to support model evaluation task
Part of https://github.com/wangkuiyi/elasticdl/issues/384. 

Extend `Worker` to support the following:
    * `distributed_evaluate()` that contains the main logic for model evaluation.
    * `report_task_result()` that reports evaluation task result (e.g. task id and error message) back to master through RPC call.
    * `report_evaluation_metrics()` that reports the computed evaluation metrics (e.g. accuracy, precision, recall, etc.) back to master through RPC call.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/master/servicer.py b/elasticdl/python/elasticdl/master/servicer.py
--- a/elasticdl/python/elasticdl/master/servicer.py
+++ b/elasticdl/python/elasticdl/master/servicer.py
@@ -214,8 +214,9 @@ def ReportEvaluationMetrics(self, request, _):
 
         with self._lock:
             for k, v in request.evaluation_metrics.items():
-                arr = tensor_to_ndarray(v)
-                self._evaluation_metrics[k] = arr
+                if v.dim:
+                    arr = tensor_to_ndarray(v)
+                    self._evaluation_metrics[k] = arr
 
             self._update_model_version()
         res.accepted = True
diff --git a/elasticdl/python/elasticdl/worker/worker.py b/elasticdl/python/elasticdl/worker/worker.py
--- a/elasticdl/python/elasticdl/worker/worker.py
+++ b/elasticdl/python/elasticdl/worker/worker.py
@@ -1,11 +1,15 @@
 import logging
 import traceback
+import time
+
 import tensorflow as tf
 assert tf.executing_eagerly()
 
 import recordio
+import numpy as np
 
 from contextlib import closing
+from collections import defaultdict
 from tensorflow.python.ops import math_ops
 from elasticdl.proto import elasticdl_pb2_grpc
 from elasticdl.proto import elasticdl_pb2
@@ -14,8 +18,12 @@
 from elasticdl.python.data.codec import TFExampleCodec
 from elasticdl.python.data.codec import BytesCodec
 
+# TODO: Move these args to train/evaluate specific functions
 # the default max number of a minibatch retrain as its gradients are not accepted by master.
 DEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64
+# the default max number of a minibatch retry as its calculated evaluation metrics are not accepted by master.
+DEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_NUM = 10
+
 
 class Worker(object):
     """ElasticDL worker"""
@@ -25,6 +33,7 @@ def __init__(self,
                  model_file,
                  channel=None,
                  max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM,
+                 max_evaluate_retry_num=DEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_NUM,
                  codec_type=None):
         """
         Arguments:
@@ -41,6 +50,7 @@ def __init__(self,
         self._input_fn = model_module.input_fn 
         self._opt_fn = model_module.optimizer
         self._loss = model_module.loss
+        self._eval_metrics_fn = model_module.eval_metrics_fn
         all_columns = self._feature_columns + model_module.label_columns()
         if codec_type == "tf_example":
             self._codec = TFExampleCodec(all_columns)
@@ -54,6 +64,7 @@ def __init__(self,
         else:
             self._stub = elasticdl_pb2_grpc.MasterStub(channel)
         self._max_retrain_num = max_retrain_num
+        self._max_evaluate_retry_num = max_evaluate_retry_num
         self._model_version = -1
         self._codec_type = codec_type
 
@@ -101,6 +112,18 @@ def report_gradient(self, grads):
         res = self._stub.ReportGradient(req)
         return res.accepted, res.model_version
 
+    def report_evaluation_metrics(self, evaluation_metrics):
+        """
+        report evaluation metrics to ps, return (accepted, model_version) from rpc call.
+        """
+        req = elasticdl_pb2.ReportEvaluationMetricsRequest()
+        for k, v in evaluation_metrics.items():
+            req.evaluation_metrics[k].CopyFrom(
+                ndarray_to_tensor(v))
+        req.model_version = self._model_version
+        res = self._stub.ReportEvaluationMetrics(req)
+        return res.accepted, res.model_version
+
     @staticmethod
     def _get_batch(reader, batch_size, decode):
         res = []
@@ -111,6 +134,13 @@ def _get_batch(reader, batch_size, decode):
             res.append(decode(record))
         return res
 
+    def _get_features_and_labels_from_record(self, record_buf):
+        batch_input_data, batch_labels = self._input_fn(record_buf)
+        features = [batch_input_data[f_col.key] for f_col in self._feature_columns]
+        if len(features) == 1:
+            features = features[0]
+        return features, batch_labels
+
     def distributed_train(self):
         """
         Distributed training.
@@ -135,16 +165,11 @@ def distributed_train(self):
                             self.get_model(
                                 max(self._model_version, min_model_version))
 
-                            batch_input_data, batch_label = self._input_fn(record_buf)
+                            features, labels = self._get_features_and_labels_from_record(record_buf)
 
                             with tf.GradientTape() as tape:
-                                inputs = []
-                                for f_col in self._feature_columns:
-                                    inputs.append(batch_input_data[f_col.key])
-                                if len(inputs) == 1:
-                                    inputs = inputs[0]
-                                outputs = self._model.call(inputs, training=True)
-                                loss = self._loss(outputs, batch_label.flatten())
+                                outputs = self._model.call(features, training=True)
+                                loss = self._loss(outputs, labels)
 
                                 # TODO:  Add regularization loss if any,
                                 #        which should be divided by the number of contributing workers.
@@ -159,7 +184,7 @@ def distributed_train(self):
                         else:
                             # Worker got stuck, fail the task.
                             # TODO: stop the worker if it fails to make any progress for some time.
-                            raise RuntimeError("Worker got stuck")
+                            raise RuntimeError("Worker got stuck during training")
             except Exception as ex:
                 err_msg = str(ex)
                 traceback.print_exc()
@@ -182,15 +207,10 @@ def local_train(self, file_list, batch_size, epoch=1, kwargs=None):
                         if not record_buf:
                             break
 
-                        data, labels = self._input_fn(record_buf)
+                        features, labels = self._get_features_and_labels_from_record(record_buf)
 
                         with tf.GradientTape() as tape:
-                            inputs = []
-                            for f_col in self._feature_columns:
-                                inputs.append(data[f_col.key])
-                            if len(inputs) == 1:
-                                inputs = inputs[0]
-                            outputs = self._model.call(inputs, training=True)
+                            outputs = self._model.call(features, training=True)
                             loss = self._loss(outputs, labels)
 
                             # Add regularization loss if any.
@@ -204,3 +224,57 @@ def local_train(self, file_list, batch_size, epoch=1, kwargs=None):
                         optimizer.apply_gradients(
                             zip(grads, self._model.trainable_variables))
                         self._logger.info("Loss is %f" % loss.numpy())
+
+    def distributed_evaluate(self, steps=None):
+        """
+        Distributed model evaluation.
+
+        Arguments:
+            steps: Evaluate the model by this many number of steps where the model is evaluated on one batch of samples
+                for each step. If `None`, evaluation will continue until reaching the end of input.
+        """
+        while True:
+            task = self.get_task()
+            if not task.shard_file_name:
+                break
+            batch_size = task.minibatch_size
+            err_msg = ""
+            try:
+                with closing(
+                        recordio.Scanner(task.shard_file_name, task.start, task.end - task.start)
+                ) as reader:
+                    min_model_version = task.model_version
+                    current_step = 0
+                    evaluation_metrics_collection = defaultdict(list)
+                    while True:
+                        current_step += 1
+                        if steps and current_step > steps:
+                            break
+                        record_buf = self._get_batch(reader, batch_size, self._codec.decode)
+                        if not record_buf:
+                            break
+
+                        for _ in range(self._max_evaluate_retry_num):
+                            self.get_model(max(self._model_version, min_model_version))
+                            features, labels = self._get_features_and_labels_from_record(record_buf)
+                            outputs = self._model.call(features, training=False)
+                            evaluation_metrics = self._eval_metrics_fn(outputs, labels)
+
+                            for k, v in evaluation_metrics.items():
+                                v_np = v.numpy()
+                                if v_np.size != 1:
+                                    raise Exception("Only metric result of length 1 is supported currently")
+                                evaluation_metrics_collection[k].append(v_np)
+
+                            evaluation_metrics = {k: np.mean(v) for k, v in evaluation_metrics_collection.items()}
+                            accepted, min_model_version = self.report_evaluation_metrics(evaluation_metrics)
+
+                            if accepted:
+                                self._logger.info("Evaluation metrics: %s" % evaluation_metrics)
+                                break
+                        else:
+                            raise RuntimeError("Worker got stuck during model evaluation")
+            except Exception as ex:
+                err_msg = str(ex)
+                traceback.print_exc()
+            self.report_task_result(task.task_id, err_msg)


Test patch:
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
@@ -10,7 +10,7 @@
 
 
 def loss(predictions, labels):
-    return tf.reduce_mean(tf.square(predictions - labels)) 
+    return tf.reduce_mean(tf.square(predictions - labels))
 
 
 def feature_columns():
@@ -37,5 +37,12 @@ def input_fn(records):
     ys = np.reshape(xs, (batch_size, 1))
     return {'x': xs}, ys
 
+
 def optimizer(lr=0.1):
     return tf.train.GradientDescentOptimizer(lr)
+
+
+def eval_metrics_fn(predictions, labels):
+    return {
+        'mse': tf.reduce_mean(tf.square(predictions - labels)),
+    }
diff --git a/elasticdl/python/tests/worker_test.py b/elasticdl/python/tests/worker_test.py
--- a/elasticdl/python/tests/worker_test.py
+++ b/elasticdl/python/tests/worker_test.py
@@ -64,9 +64,9 @@ def test_local_train_bytes(self):
     def test_local_train_tf_example(self):
         self.local_train("tf_example")
 
-    def distributed_train(self, codec_type):
+    def distributed_train_and_evaluate(self, codec_type, training=True):
         """
-        Run Worker.distributed_train with a local master.
+        Run distributed training and evaluation with a local master.
         grpc calls are mocked by local master call.
         """
 
@@ -83,6 +83,13 @@ def mock_ReportGradient(req):
                 master._version += 1
             return master.ReportGradient(req, None)
 
+        def mock_ReportEvaluationMetrics(req):
+            if 2 < master._version < 80:
+                # For testing of evaluation retries when evaluation metrics are not accepted.
+                # Increase master version so the evaluation metrics will not be accepted.
+                master._version += 1
+            return master.ReportEvaluationMetrics(req, None)
+
         def mock_ReportTaskResult(req):
             return master.ReportTaskResult(req, None)
 
@@ -111,11 +118,16 @@ def mock_ReportTaskResult(req):
             worker._stub, "GetModel", mock_GetModel
         ), mock.patch.object(
             worker._stub, "ReportGradient", mock_ReportGradient
+        ), mock.patch.object(
+            worker._stub, 'ReportEvaluationMetrics', mock_ReportEvaluationMetrics
         ), mock.patch.object(
             worker._stub, "ReportTaskResult", mock_ReportTaskResult
         ):
             try:
-                worker.distributed_train()
+                if training:
+                    worker.distributed_train()
+                else:
+                    worker.distributed_evaluate(steps=2)
                 res = True
             except Exception as ex:
                 print(ex)
@@ -129,10 +141,16 @@ def mock_ReportTaskResult(req):
         self.assertTrue(not task.shard_file_name)
 
     def test_distributed_train_bytes(self):
-        self.distributed_train("bytes")
+        self.distributed_train_and_evaluate("bytes", training=True)
+
+    def test_distributed_evaluate_bytes(self):
+        self.distributed_train_and_evaluate("bytes", training=False)
 
     def test_distributed_train_tf_example(self):
-        self.distributed_train("tf_example")
+        self.distributed_train_and_evaluate("tf_example", training=True)
+
+    def test_distributed_evaluate_tf_example(self):
+        self.distributed_train_and_evaluate("tf_example", training=False)
 
 
 if __name__ == "__main__":