Fix should validate callback train end (#5542)

* Add on_end method to ShouldValidateCallback

* Add test for _should_validate_this_epoch at end of training

* Update changelog

* Use epoch argument if provided

CHANGELOG.md

@@ -26,8 +26,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Removed a spurious error message "'torch.cuda' has no attribute '_check_driver'" that would be appear in the logs
   when a `ConfigurationError` for missing GPU was raised.
 - Load model on CPU post training to save GPU memory.
-- Fixed a bug in `ShouldValidateCallback` that leads to valuation occuring after the first epoch regardless of `validation_start` value.
-- Fixed a bug in `ShouldValidateCallback` that leads to valuation occuring every `validation_interval + 1` epochs, instead of every `validation_interval` epochs.
+- Fixed a bug in `ShouldValidateCallback` that leads to validation occuring after the first epoch regardless of `validation_start` value.
+- Fixed a bug in `ShouldValidateCallback` that leads to validation occuring every `validation_interval + 1` epochs, instead of every `validation_interval` epochs.
+- Fixed a bug in `ShouldValidateCallback` that leads to validation never occuring at the end of training.
 
 ### Removed
 

allennlp/training/callbacks/should_validate.py

@@ -40,6 +40,17 @@ def on_epoch(
     ) -> None:
         trainer._should_validate_this_epoch = self._should_validate(epoch=epoch + 1)
 
+    def on_end(
+        self,
+        trainer: "GradientDescentTrainer",
+        metrics: Dict[str, Any] = None,
+        epoch: int = None,
+        is_primary: bool = True,
+        **kwargs,
+    ) -> None:
+        epoch = epoch + 1 if epoch is not None else trainer._epochs_completed
+        trainer._should_validate_this_epoch = self._should_validate(epoch=epoch)
+
     def _should_validate(self, epoch: int) -> bool:
         should_validate = True
         if self._validation_start is not None and epoch < self._validation_start:

tests/training/trainer_test.py

@@ -1374,7 +1374,11 @@ def test_should_validate_callback(self):
         assert not trainer._should_validate_this_epoch
 
         # Satisfies both 'validation_start' and 'validation_interval'
-        callback.on_epoch(trainer, metrics={}, epoch=5)
+        callback.on_epoch(trainer, metrics={}, epoch=3)
+        assert trainer._should_validate_this_epoch
+
+        # Check that final validation happens on the last epoch
+        callback.on_end(trainer)
         assert trainer._should_validate_this_epoch