HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fix ClearML Mosaic callback to 'on_train_epoch_end' (#92)

Browse source
This commit is contained in:
Glenn Jocher 2022-12-25 17:04:24 +01:00 • committed by GitHub
parent 0298821467
commit 249dfbdc05
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 17 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

12
ultralytics/yolo/utils/callbacks/clearml.py
View file

@ -1,6 +1,3 @@
import os
from pathlib import Path
from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
try:
@ -30,12 +27,9 @@ def on_train_start(trainer):
task.connect(dict(trainer.args), name='General')
def on_epoch_start(trainer):
def on_train_epoch_end(trainer):
if trainer.epoch == 1:
plots = [filename for filename in os.listdir(trainer.save_dir) if filename.startswith("train_batch")]
imgs_dict = {f"train_batch_{i}": Path(trainer.save_dir) / img for i, img in enumerate(plots)}
if imgs_dict:
_log_images(imgs_dict, "Mosaic", trainer.epoch)
_log_images({f.stem: str(f) for f in trainer.save_dir.glob('train_batch*.jpg')}, "Mosaic", trainer.epoch)
def on_val_end(trainer):
@ -55,6 +49,6 @@ def on_train_end(trainer):
callbacks = {
"on_train_start": on_train_start,
"on_epoch_start": on_epoch_start,
"on_train_epoch_end": on_train_epoch_end,
"on_val_end": on_val_end,
"on_train_end": on_train_end} if clearml else {}