Skip to content

Commit

Permalink
Add stop_training=False flag to callbacks (#6365)
Browse files Browse the repository at this point in the history
* 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 <glenn.jocher@ultralytics.com>
  • Loading branch information
haimat and glenn-jocher authored Jan 23, 2022
1 parent c43439a commit 8efe977
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 0 deletions.
2 changes: 2 additions & 0 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions utils/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down

0 comments on commit 8efe977

Please sign in to comment.