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

@ -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 {}