- added callback after_model_and_loss_is_applied (for #152)

inferno-pytorch · Nov 3, 2018 · 2ef65f6 · 2ef65f6
1 parent 45220b4
commit 2ef65f6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/inferno/trainers/basic.py b/inferno/trainers/basic.py
@@ -1426,6 +1426,8 @@ def train_for(self, num_iterations=None, break_callback=None):
                 # Apply model, compute loss and backprop
                 prediction, loss = self.apply_model_and_loss(inputs, target, backward=True,
                                                              mode='train')
+            self.callbacks.call(self.callbacks.AFTER_MODEL_AND_LOSS_IS_APPLIED,
+                                prediction=prediction, loss=loss, iteration_num=iteration_num)
             # Compute metric
             if self.metric_is_defined and self.evaluate_metric_now:
                 self._last_metric_evaluated_at_epoch = self._epoch_count

diff --git a/inferno/trainers/callbacks/base.py b/inferno/trainers/callbacks/base.py
@@ -18,6 +18,7 @@ class CallbackEngine(object):
     BEGIN_OF_EPOCH = 'begin_of_epoch'
     END_OF_EPOCH = 'end_of_epoch'
     BEGIN_OF_TRAINING_ITERATION = 'begin_of_training_iteration'
+    AFTER_MODEL_AND_LOSS_IS_APPLIED = 'after_model_and_loss_is_applied'
     END_OF_TRAINING_ITERATION = 'end_of_training_iteration'
     BEGIN_OF_VALIDATION_RUN = 'begin_of_validation_run'
     END_OF_VALIDATION_RUN = 'end_of_validation_run'
@@ -33,6 +34,7 @@ class CallbackEngine(object):
                 BEGIN_OF_EPOCH,
                 END_OF_EPOCH,
                 BEGIN_OF_TRAINING_ITERATION,
+                AFTER_MODEL_AND_LOSS_IS_APPLIED,
                 END_OF_TRAINING_ITERATION,
                 BEGIN_OF_VALIDATION_RUN,
                 END_OF_VALIDATION_RUN,