Add validation criterion callback hook (#235)

* Add validation criterion callback hook

* Fix bug in decorator tests

tests/callbacks/test_callbacks.py

@@ -24,6 +24,7 @@ def test_empty_methods(self):
         self.assertIsNone(callback.on_forward_validation({}))
         self.assertIsNone(callback.on_end_validation({}))
         self.assertIsNone(callback.on_step_validation({}))
+        self.assertIsNone(callback.on_criterion_validation({}))
 
 
 class TestCallbackList(TestCase):

tests/callbacks/test_decorators.py

@@ -90,6 +90,12 @@ def example(state):
         state = 'test'
         self.assertTrue(callbacks.on_forward_validation(example).on_forward_validation(state) == state)
 
+    def test_on_criterion_validation(self):
+        def example(state):
+            return state
+        state = 'test'
+        self.assertTrue(callbacks.on_criterion_validation(example).on_criterion_validation(state) == state)
+
     def test_on_end_validation(self):
         def example(state):
             return state
@@ -108,7 +114,7 @@ def example(state):
         state = {'test': 'test', torchbearer.LOSS: 0}
         callbacks.add_to_loss(example).on_criterion(state)
         self.assertTrue(state[torchbearer.LOSS] == 1)
-        callbacks.add_to_loss(example).on_step_validation(state)
+        callbacks.add_to_loss(example).on_criterion_validation(state)
         self.assertTrue(state[torchbearer.LOSS] == 2)
 
 

torchbearer/callbacks/callbacks.py

@@ -134,6 +134,16 @@ def on_forward_validation(self, state):
         """
         pass
 
+    def on_criterion_validation(self, state):
+        """Perform some action with the given state as context after the criterion evaluation has been completed
+        with the validation data.
+
+        :param state: The current state dict of the :class:`Model`.
+        :type state: dict[str,any]
+
+        """
+        pass
+
     def on_end_validation(self, state):
         """Perform some action with the given state as context at the end of the validation loop.
 
@@ -306,6 +316,15 @@ def on_forward_validation(self, state):
         """
         self._for_list(lambda callback: callback.on_forward_validation(state))
 
+    def on_criterion_validation(self, state):
+        """Call on_criterion_validation on each callback in turn with the given state.
+
+        :param state: The current state dict of the :class:`Model`.
+        :type state: dict[str,any]
+
+        """
+        self._for_list(lambda callback: callback.on_criterion_validation(state))
+
     def on_end_validation(self, state):
         """Call on_end_validation on each callback in turn with the given state.
 

torchbearer/callbacks/decorators.py

@@ -197,6 +197,20 @@ def on_forward_validation(func):
     return new_callback
 
 
+def on_criterion_validation(func):
+    """ The :func:`on_criterion_validation` decorator is used to initialise a :class:`.Callback` with :meth:`~.Callback.on_criterion_validation`
+    calling the decorated function
+
+    :param func: The function(state) to *decorate*
+    :type func: function
+    :return: Initialised callback with :meth:`.Callback.on_criterion_validation` calling func
+    :rtype: :class:`.Callback`
+    """
+    new_callback = Callback()
+    new_callback.on_criterion_validation = func
+    return new_callback
+
+
 def on_end_validation(func):
     """ The :func:`on_end_validation` decorator is used to initialise a :class:`.Callback` with :meth:`~.Callback.on_end_validation`
     calling the decorated function
@@ -239,5 +253,5 @@ def add_to_loss_func(state):
 
     new_callback = Callback()
     new_callback.on_criterion = add_to_loss_func
-    new_callback.on_step_validation = add_to_loss_func
+    new_callback.on_criterion_validation = add_to_loss_func
     return new_callback

torchbearer/torchbearer.py

@@ -253,6 +253,7 @@ def _test_loop(self, state, callbacks, pass_state, batch_loader, num_steps=None)
                 # Loss and metrics
                 if torchbearer.Y_TRUE in state:
                     state[torchbearer.LOSS] = state[torchbearer.CRITERION](state[torchbearer.Y_PRED], state[torchbearer.Y_TRUE])
+                    callbacks.on_criterion_validation(state)
                     state[torchbearer.METRICS] = state[torchbearer.METRIC_LIST].process(state)
 
                 callbacks.on_step_validation(state)