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:
@ -18,7 +18,7 @@ def _log_images(imgs_dict, group="", step=0):
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
# TODO: reuse existing task
|
||||
task = Task.init(project_name=trainer.args.project if trainer.args.project != 'runs/train' else 'YOLOv8',
|
||||
task = Task.init(project_name=trainer.args.project or "YOLOv8",
|
||||
task_name=trainer.args.name,
|
||||
tags=['YOLOv8'],
|
||||
output_uri=True,
|
||||
@ -32,7 +32,7 @@ def on_train_epoch_end(trainer):
|
||||
_log_images({f.stem: str(f) for f in trainer.save_dir.glob('train_batch*.jpg')}, "Mosaic", trainer.epoch)
|
||||
|
||||
|
||||
def on_val_end(trainer):
|
||||
def on_fit_epoch_end(trainer):
|
||||
if trainer.epoch == 0:
|
||||
model_info = {
|
||||
"Parameters": get_num_params(trainer.model),
|
||||
@ -50,5 +50,5 @@ def on_train_end(trainer):
|
||||
callbacks = {
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_train_epoch_end": on_train_epoch_end,
|
||||
"on_val_end": on_val_end,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_train_end": on_train_end} if clearml else {}
|
||||
|
||||
Reference in New Issue
Block a user