Add `stop_training=False` flag to callbacks (#6365)

* New flag 'stop_training' in util.callbacks.Callbacks class to prematurely stop training from callback handler

* Removed most of the new checks, leaving only the one after calling 'on_train_batch_end'

* Cleanup

Co-authored-by: Glenn Jocher <[email protected]>

train.py

@@ -352,6 +352,8 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
                 pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % (
                     f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
                 callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn)
+                if callbacks.stop_training:
+                    return
             # end batch ------------------------------------------------------------------------------------------------
 
         # Scheduler

utils/callbacks.py

@@ -35,6 +35,7 @@ class Callbacks:
             'on_params_update': [],
             'teardown': [],
         }
+        self.stop_training = False  # set True to interrupt training
 
     def register_action(self, hook, name='', callback=None):
         """