From 8efe97719c9b3b77c9db9b5c8592e051b7f0c9a7 Mon Sep 17 00:00:00 2001 From: Matthias Date: Sun, 23 Jan 2022 03:37:21 +0100 Subject: [PATCH] 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 --- train.py | 2 ++ utils/callbacks.py | 1 + 2 files changed, 3 insertions(+) diff --git a/train.py b/train.py index b20b7dbb2..510377e11 100644 --- 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 index 13d82ebc2..c51c268f2 100644 --- a/utils/callbacks.py +++ b/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): """