Develop an API to get training epoch (#2488)

* Check whether to register hooks according to HOROVOD_ELASTIC

* Develop an API to get training epoch

* Register hooks

* Add unittest

* Fic by comments

* Fix unittest

Author: workingloong

elasticai_api/common/data_shard_service.py

@@ -81,6 +81,9 @@ def _report_training_params(self):
                 num_minibatches_per_shard=self._num_minibatches_per_shard,
             )
 
+    def get_minibatch_count_per_epoch(self):
+        return self._dataset_size // self._batch_size
+
     def get_current_task(self):
         return self._current_task
 

elasticai_api/pytorch/controller.py

@@ -107,6 +107,15 @@ def __init__(self, master_client, data_shard_service):
             os.getenv(WorkerEnv.WORKER_NUM, 1)
         )
         self.global_completed_batch_num = 0
+        self.batch_count_per_epoch = (
+            self.data_shard_service.get_minibatch_count_per_epoch()
+        )
+
+    def get_current_epoch(self):
+        return self.global_completed_batch_num // self.batch_count_per_epoch
+
+    def set_resume_epoch(self, epoch):
+        self.global_completed_batch_num = epoch * self.batch_count_per_epoch
 
     def set_broadcast_model(self, model):
         self._model = model
@@ -175,8 +184,8 @@ def reset_backward_passes_per_step(self):
         ):
             world_size = hvd.size()
             rank = hvd.rank()
-            self.backward_passes_per_step = int(
-                self.global_batch_num_per_step / world_size
+            self.backward_passes_per_step = (
+                self.global_batch_num_per_step // world_size
             )
             if rank < self.global_batch_num_per_step % world_size:
                 self.backward_passes_per_step += 1

elasticdl/python/tests/allreduce_trainer_test.py

@@ -100,7 +100,7 @@ def test_elastic_run(self):
                 rendezvous_id=1, rank_id=0, world_size=1, rendezvous_port=0
             )
         )
-        data_shard_service = DataShardService(master_client, 1)
+        data_shard_service = DataShardService(master_client, 1, 1, 10)
         controller = AllReduceController(master_client, data_shard_service)
         elastic_run = controller.elastic_run(self.train)
         elastic_run()
@@ -115,7 +115,7 @@ def setUp(self):
                 rendezvous_id=1, rank_id=0, world_size=1, rendezvous_port=0
             )
         )
-        data_shard_service = DataShardService(master_client, 1)
+        data_shard_service = DataShardService(master_client, 1, 1, 10)
         self.controller = TensorFlowV2AllReduceController(
             master_client, data_shard_service
         )
@@ -145,7 +145,7 @@ def setUp(self):
                 rendezvous_id=1, rank_id=0, world_size=1, rendezvous_port=0
             )
         )
-        data_shard_service = DataShardService(master_client, 1)
+        data_shard_service = DataShardService(master_client, 1, 1, 10)
         self.controller = PyTorchAllReduceController(
             master_client, data_shard_service
         )
@@ -171,7 +171,7 @@ def test_elastic_run(self):
         self.assertEqual(self.controller.global_completed_batch_num, 1)
 
     def test_create_elastic_controller(self):
-        controller = create_elastic_controller(batch_size=64)
+        controller = create_elastic_controller(batch_size=64, dataset_size=128)
         self.assertIsNotNone(controller)
         self.assertIsNotNone(controller.data_shard_service._mc)
         self.assertEqual(controller.data_shard_service._batch_size, 64)
@@ -187,6 +187,14 @@ def test_reset_backward_passes_per_step(self):
         self.controller.reset_backward_passes_per_step()
         self.assertEqual(self.controller.backward_passes_per_step, 2)
 
+    def test_get_epoch(self):
+        self.controller.batch_count_per_epoch = 10
+        self.controller.global_completed_batch_num = 78
+        epoch = self.controller.get_current_epoch()
+        self.assertEqual(epoch, 7)
+        self.controller.set_resume_epoch(5)
+        self.assertEqual(self.controller.global_completed_batch_num, 50)
+
 
 if __name__ == "__main__":
     unittest.main()

model_zoo/mnist/mnist_pytorch.py

@@ -131,7 +131,6 @@ def train(args):
     device = torch.device("cuda" if use_cuda else "cpu")
     train_data = torchvision.datasets.ImageFolder(args.training_data)
     test_data = torchvision.datasets.ImageFolder(args.validation_data)
-    batch_num_per_epoch = int(len(train_data.imgs) / args.batch_size)
 
     allreduce_controller = create_elastic_controller(
         batch_size=args.batch_size,
@@ -171,10 +170,7 @@ def train(args):
             data, target = data.to(device), target.to(device)
             loss = elastic_train_one_batch(model, optimizer, data, target)
             print("loss = {}, step = {}".format(loss, batch_idx))
-            new_epoch = int(
-                allreduce_controller.global_completed_batch_num
-                / batch_num_per_epoch
-            )
+            new_epoch = allreduce_controller.get_current_epoch()
             if new_epoch > epoch:
                 epoch = new_epoch
                 # Set epoch of the scheduler