Repository: sql-machine-learning/elasticdl

Issue:
Support initializing the training model from a checkpoint 
When submitting an ElasticDL job, the user can optionally add `--init_from_checkpoint=checkpoint_file` to initialize the training model from `checkpoint_file`.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/master/main.py b/elasticdl/python/elasticdl/master/main.py
--- a/elasticdl/python/elasticdl/master/main.py
+++ b/elasticdl/python/elasticdl/master/main.py
@@ -79,6 +79,11 @@ def _parse_args():
         help="the number of workers used in training",
         default=0,
     )
+    parser.add_argument(
+        "--init_from_checkpoint",
+        help="The checkpoint file to initialize the training model",
+        default="",
+    )
     parser.add_argument(
         "--checkpoint_dir",
         help="The directory to store the checkpoint files",
@@ -183,6 +188,7 @@ def main():
             optimizer,
             task_q,
             init_var=model_inst.trainable_variables,
+            init_from_checkpoint=args.init_from_checkpoint,
             checkpoint_dir=args.checkpoint_dir,
             checkpoint_steps=args.checkpoint_steps,
             keep_checkpoint_max=args.keep_checkpoint_max,
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
@@ -26,6 +26,7 @@ def __init__(
         task_q,
         *,
         init_var,
+        init_from_checkpoint,
         checkpoint_dir,
         checkpoint_steps,
         keep_checkpoint_max
@@ -47,6 +48,8 @@ def __init__(
         self._evaluation_metrics = {}
         for var in init_var:
             self.set_model_var(var.name, var.numpy())
+        if init_from_checkpoint:
+            self.load_checkpoint_file(init_from_checkpoint)
         self._checkpoint_dir = checkpoint_dir
         self._checkpoint_steps = checkpoint_steps
         self._keep_checkpoint_max = keep_checkpoint_max


Test patch:
diff --git a/elasticdl/python/tests/checkpoint_test.py b/elasticdl/python/tests/checkpoint_test.py
--- a/elasticdl/python/tests/checkpoint_test.py
+++ b/elasticdl/python/tests/checkpoint_test.py
@@ -16,7 +16,10 @@
 from elasticdl.python.elasticdl.worker.worker import Worker
 from elasticdl.python.elasticdl.common.model_helper import load_user_model
 from elasticdl.python.elasticdl.master.task_queue import _TaskQueue
-from elasticdl.python.elasticdl.common.model_helper import save_checkpoint_to_file, load_from_checkpoint_file
+from elasticdl.python.elasticdl.common.model_helper import (
+    save_checkpoint_to_file,
+    load_from_checkpoint_file,
+)
 from elasticdl.proto import elasticdl_pb2
 from elasticdl.python.data.codec import BytesCodec, TFExampleCodec
 
@@ -46,18 +49,19 @@ def create_recordio_file(size, codec_type):
 
 class CheckpointTest(unittest.TestCase):
     def testSaveLoadCheckpoint(self):
+        init_var = m.model.trainable_variables
         master = MasterServicer(
-            2, 3, None, None,
-            init_var=[],
+            2,
+            3,
+            None,
+            None,
+            init_var=init_var,
+            init_from_checkpoint="",
             checkpoint_dir="",
             checkpoint_steps=0,
-            keep_checkpoint_max=0
+            keep_checkpoint_max=0,
         )
 
-        model_inst = m.model
-        for variable in model_inst.trainable_variables:
-            master.set_model_var(variable.name, variable.numpy())
-
         req = elasticdl_pb2.GetModelRequest()
         req.min_version = 0
         model = master.GetModel(req, None)
@@ -71,6 +75,42 @@ def testSaveLoadCheckpoint(self):
         for k in model.param:
             self.assertEqual(model.param[k], pb_model.param[k])
 
+    def testInitCheckpoint(self):
+        init_var = m.model.trainable_variables
+        req = elasticdl_pb2.GetModelRequest()
+        req.min_version = 0
+
+        master = MasterServicer(
+            2,
+            3,
+            None,
+            None,
+            init_var=init_var,
+            init_from_checkpoint="",
+            checkpoint_dir="",
+            checkpoint_steps=0,
+            keep_checkpoint_max=0,
+        )
+        model = master.GetModel(req, None)
+
+        tmp_file = tempfile.NamedTemporaryFile()
+        save_checkpoint_to_file(model, tmp_file.name)
+
+        master2 = MasterServicer(
+            2,
+            3,
+            None,
+            None,
+            init_var=init_var,
+            init_from_checkpoint=tmp_file.name,
+            checkpoint_dir="",
+            checkpoint_steps=0,
+            keep_checkpoint_max=0,
+        )
+        model2 = master2.GetModel(req, None)
+
+        self.assertEqual(model, model2)
+
     def testCheckpintArguments(self):
         """
         Run Worker.distributed_train with a local master.
@@ -91,12 +131,7 @@ def mock_ReportTaskResult(req):
 
         codec_type = "bytes"
         channel = grpc.insecure_channel("localhost:9999")
-        worker = Worker(
-            1,
-            _module_file,
-            channel=channel,
-            codec_type=codec_type,
-        )
+        worker = Worker(1, _module_file, channel=channel, codec_type=codec_type)
 
         # save checkpoint file every 2 steps
         # keep at most 5 recent checkpoint files
@@ -111,18 +146,20 @@ def mock_ReportTaskResult(req):
             num_epoch=1,
             task_type=elasticdl_pb2.TRAINING,
         )
-        master = MasterServicer(2,
-                                2,
-                                worker._opt_fn(),
-                                task_q,
-                                init_var=[],
-                                checkpoint_dir=checkpoint_dir,
-                                checkpoint_steps=checkpoint_steps,
-                                keep_checkpoint_max=keep_checkpoint_max
-                                )
-
-        for var in worker._model.trainable_variables:
-            master.set_model_var(var.name, var.numpy())
+        master = MasterServicer(
+            2,
+            2,
+            worker._opt_fn(),
+            task_q,
+            init_var=worker._model.trainable_variables,
+            init_from_checkpoint="",
+            checkpoint_dir=checkpoint_dir,
+            checkpoint_steps=checkpoint_steps,
+            keep_checkpoint_max=keep_checkpoint_max,
+        )
+
+        # for var in worker._model.trainable_variables:
+        #    master.set_model_var(var.name, var.numpy())
 
         with mock.patch.object(
             worker._stub, "GetTask", mock_GetTask
@@ -142,8 +179,16 @@ def mock_ReportTaskResult(req):
 
         self.assertTrue(res)
         checkpoint_files = sorted(os.listdir(checkpoint_dir))
-        self.assertEqual(checkpoint_files,
-                         ['model_v24.chkpt', 'model_v26.chkpt', 'model_v28.chkpt', 'model_v30.chkpt', 'model_v32.chkpt'])
+        self.assertEqual(
+            checkpoint_files,
+            [
+                "model_v24.chkpt",
+                "model_v26.chkpt",
+                "model_v28.chkpt",
+                "model_v30.chkpt",
+                "model_v32.chkpt",
+            ],
+        )
 
 
 if __name__ == "__main__":
diff --git a/elasticdl/python/tests/servicer_test.py b/elasticdl/python/tests/servicer_test.py
--- a/elasticdl/python/tests/servicer_test.py
+++ b/elasticdl/python/tests/servicer_test.py
@@ -49,6 +49,7 @@ def testGetEmptyTask(self):
                 task_type=elasticdl_pb2.TRAINING
             ),
             init_var=[],
+            init_from_checkpoint="",
             checkpoint_dir="",
             checkpoint_steps=0,
             keep_checkpoint_max=0
@@ -70,6 +71,7 @@ def testGetEmptyTask(self):
     def testGetModel(self):
         master = MasterServicer(2, 3, None, None,
                                 init_var=[],
+                                init_from_checkpoint="",
                                 checkpoint_dir="",
                                 checkpoint_steps=0,
                                 keep_checkpoint_max=0
@@ -123,6 +125,7 @@ def makeGrad():
             tf.train.GradientDescentOptimizer(0.1),
             None,
             init_var=[],
+            init_from_checkpoint="",
             checkpoint_dir="",
             checkpoint_steps=0,
             keep_checkpoint_max=0
@@ -213,6 +216,7 @@ def makeEvaluationMetrics():
             tf.train.GradientDescentOptimizer(0.1),
             None,
             init_var=[],
+            init_from_checkpoint="",
             checkpoint_dir="",
             checkpoint_steps=0,
             keep_checkpoint_max=0
@@ -250,6 +254,7 @@ def testReportTaskResult(self):
         )
         master = MasterServicer(3, 3, None, task_q,
                                 init_var=[],
+                                init_from_checkpoint="",
                                 checkpoint_dir="",
                                 checkpoint_steps=0,
                                 keep_checkpoint_max=0
@@ -289,6 +294,7 @@ def testReportTaskResult(self):
     def testUserDefinedModel(self):
         master = MasterServicer(2, 3, None, None,
                                 init_var=[],
+                                init_from_checkpoint="",
                                 checkpoint_dir="",
                                 checkpoint_steps=0,
                                 keep_checkpoint_max=0
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
@@ -86,6 +86,7 @@ def mock_ReportTaskResult(req):
         task_q = _TaskQueue({filename: 128}, record_per_task=64, num_epoch=1, task_type=task_type)
         master = MasterServicer(2, 16, worker._opt_fn(), task_q,
                                 init_var=[],
+                                init_from_checkpoint="",
                                 checkpoint_dir="",
                                 checkpoint_steps=0,
                                 keep_checkpoint_max=0)