Add pred, export and val callbacks (#126)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia
2023-01-01 22:46:10 +05:30
committed by GitHub
parent 63c7a74691
commit c6eb6720de
8 changed files with 176 additions and 57 deletions

View File

@ -1,3 +1,7 @@
# Ultralytics YOLO base callbacks
# Trainer callbacks ----------------------------------------------------------------------------------------------------
def on_pretrain_routine_start(trainer):
pass
@ -34,26 +38,6 @@ def on_train_epoch_end(trainer):
pass
def on_val_start(trainer):
pass
def on_val_batch_start(trainer):
pass
def on_val_image_end(trainer):
pass
def on_val_batch_end(trainer):
pass
def on_val_end(trainer):
pass
def on_fit_epoch_end(trainer):
pass
@ -74,7 +58,51 @@ def teardown(trainer):
pass
# Validator callbacks --------------------------------------------------------------------------------------------------
def on_val_start(validator):
pass
def on_val_batch_start(validator):
pass
def on_val_batch_end(validator):
pass
def on_val_end(validator):
pass
# Predictor callbacks --------------------------------------------------------------------------------------------------
def on_predict_start(predictor):
pass
def on_predict_batch_start(predictor):
pass
def on_predict_batch_end(predictor):
pass
def on_predict_end(predictor):
pass
# Exporter callbacks ---------------------------------------------------------------------------------------------------
def on_export_start(exporter):
pass
def on_export_end(exporter):
pass
default_callbacks = {
# Run in trainer
'on_pretrain_routine_start': on_pretrain_routine_start,
'on_pretrain_routine_end': on_pretrain_routine_end,
'on_train_start': on_train_start,
@ -84,16 +112,27 @@ default_callbacks = {
'on_before_zero_grad': on_before_zero_grad,
'on_train_batch_end': on_train_batch_end,
'on_train_epoch_end': on_train_epoch_end,
'on_val_start': on_val_start,
'on_val_batch_start': on_val_batch_start,
'on_val_image_end': on_val_image_end,
'on_val_batch_end': on_val_batch_end,
'on_val_end': on_val_end,
'on_fit_epoch_end': on_fit_epoch_end, # fit = train + val
'on_model_save': on_model_save,
'on_train_end': on_train_end,
'on_params_update': on_params_update,
'teardown': teardown}
'teardown': teardown,
# Run in validator
'on_val_start': on_val_start,
'on_val_batch_start': on_val_batch_start,
'on_val_batch_end': on_val_batch_end,
'on_val_end': on_val_end,
# Run in predictor
'on_predict_start': on_predict_start,
'on_predict_batch_start': on_predict_batch_start,
'on_predict_batch_end': on_predict_batch_end,
'on_predict_end': on_predict_end,
# Run in exporter
'on_export_start': on_export_start,
'on_export_end': on_export_end}
def add_integration_callbacks(trainer):