EpochCallBack for tracking epoch (#4540)

* EpochCallback for tracking epoch in the model

* minor lint

* updated CHANGELOG

* added unit test for track epoch callback

* Update allennlp/training/trainer.py

Co-authored-by: Matt Gardner <[email protected]>

Author: nitishgupta

CHANGELOG.md

@@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added the option to specify `requires_grad: false` within an optimizer's parameter groups.
 - Added the `file-friendly-logging` flag back to the `train` command. Also added this flag to the `predict`, `evaluate`, and `find-learning-rate` commands.
+- Added an `EpochCallback` to track current epoch as a model class member. 
 
 ### Removed
 

allennlp/training/__init__.py

@@ -1,4 +1,10 @@
 from allennlp.training.checkpointer import Checkpointer
 from allennlp.training.tensorboard_writer import TensorboardWriter
 from allennlp.training.no_op_trainer import NoOpTrainer
-from allennlp.training.trainer import Trainer, GradientDescentTrainer, BatchCallback, EpochCallback
+from allennlp.training.trainer import (
+    Trainer,
+    GradientDescentTrainer,
+    BatchCallback,
+    EpochCallback,
+    TrackEpochCallback,
+)

allennlp/training/trainer.py

@@ -176,6 +176,29 @@ def __call__(
 EpochCallback.register("null")(EpochCallback)
 
 
+@EpochCallback.register("track_epoch_callback")
+class TrackEpochCallback:
+    """
+    A callback that you can pass to the `GradientDescentTrainer` to access the current epoch number
+    in your model during training. This callback sets `model.epoch`, which can be read inside of
+    `model.forward()`. Since the EpochCallback passes `epoch=-1`
+    at the start of the training, we set `model.epoch = epoch + 1` which now denotes the number of
+    completed epochs at a given training state.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def __call__(
+        self,
+        trainer: "GradientDescentTrainer",
+        metrics: Dict[str, Any],
+        epoch: int,
+        is_master: bool,
+    ) -> None:
+        trainer.model.epoch = epoch + 1
+
+
 @Trainer.register("gradient_descent", constructor="from_partial_objects")
 class GradientDescentTrainer(Trainer):
     """

tests/training/trainer_test.py

@@ -28,6 +28,7 @@
     TensorboardWriter,
     BatchCallback,
     EpochCallback,
+    TrackEpochCallback,
 )
 from allennlp.training.learning_rate_schedulers import CosineWithRestarts
 from allennlp.training.learning_rate_schedulers import ExponentialLearningRateScheduler
@@ -986,6 +987,19 @@ def __call__(
         expected_calls = [epoch for epoch in range(-1, 4)]
         assert trainer.epoch_callback_calls == expected_calls
 
+    def test_track_epoch_callback(self):
+        num_epochs = 4
+        trainer = GradientDescentTrainer(
+            self.model,
+            self.optimizer,
+            self.data_loader,
+            num_epochs=num_epochs,
+            validation_data_loader=self.validation_data_loader,
+            epoch_callbacks=[TrackEpochCallback()],
+        )
+        trainer.train()
+        assert trainer.model.epoch == num_epochs
+
     def test_total_loss_is_average_of_batch_loss(self):
 
         batches_per_epoch = 3