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]>
ultralytics · Jan 23, 2022 · 8efe977 · 8efe977
1 parent c43439a
commit 8efe977
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/train.py b/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

diff --git a/utils/callbacks.py b/utils/callbacks.py
@@ -35,6 +35,7 @@ def __init__(self):
             'on_params_update': [],
             'teardown': [],
         }
+        self.stop_training = False  # set True to interrupt training
 
     def register_action(self, hook, name='', callback=None):
         """