From 249dfbdc05316ac2055541a797fb676699c57de5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 25 Dec 2022 17:04:24 +0100 Subject: [PATCH] Fix ClearML Mosaic callback to 'on_train_epoch_end' (#92) --- ultralytics/yolo/engine/validator.py | 1 + ultralytics/yolo/utils/callbacks/clearml.py | 12 +++--------- ultralytics/yolo/utils/metrics.py | 10 +++++----- ultralytics/yolo/v8/detect/train.py | 15 +++++++-------- ultralytics/yolo/v8/segment/train.py | 3 +-- 5 files changed, 17 insertions(+), 24 deletions(-) diff --git a/ultralytics/yolo/engine/validator.py b/ultralytics/yolo/engine/validator.py index b150614..03bf951 100644 --- a/ultralytics/yolo/engine/validator.py +++ b/ultralytics/yolo/engine/validator.py @@ -47,6 +47,7 @@ class BaseValidator: model = model.half() if self.args.half else model.float() self.model = model self.loss = torch.zeros_like(trainer.loss_items, device=trainer.device) + self.args.plots = trainer.epoch == trainer.epochs - 1 # always plot final epoch else: assert model is not None, "Either trainer or model is needed for validation" self.device = select_device(self.args.device, self.args.batch_size) diff --git a/ultralytics/yolo/utils/callbacks/clearml.py b/ultralytics/yolo/utils/callbacks/clearml.py index 400944b..3cfd4b1 100644 --- a/ultralytics/yolo/utils/callbacks/clearml.py +++ b/ultralytics/yolo/utils/callbacks/clearml.py @@ -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 {} diff --git a/ultralytics/yolo/utils/metrics.py b/ultralytics/yolo/utils/metrics.py index 30b1b72..92cedb3 100644 --- a/ultralytics/yolo/utils/metrics.py +++ b/ultralytics/yolo/utils/metrics.py @@ -343,7 +343,7 @@ def compute_ap(recall, precision): return ap, mpre, mrec -def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=(), eps=1e-16, prefix=""): +def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir=Path(), names=(), eps=1e-16, prefix=""): """ Compute the average precision, given the recall and precision curves. Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. # Arguments @@ -398,10 +398,10 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data names = dict(enumerate(names)) # to dict if plot: - plot_pr_curve(px, py, ap, Path(save_dir) / f'{prefix}PR_curve.png', names) - plot_mc_curve(px, f1, Path(save_dir) / f'{prefix}F1_curve.png', names, ylabel='F1') - plot_mc_curve(px, p, Path(save_dir) / f'{prefix}P_curve.png', names, ylabel='Precision') - plot_mc_curve(px, r, Path(save_dir) / f'{prefix}R_curve.png', names, ylabel='Recall') + plot_pr_curve(px, py, ap, save_dir / f'{prefix}PR_curve.png', names) + plot_mc_curve(px, f1, save_dir / f'{prefix}F1_curve.png', names, ylabel='F1') + plot_mc_curve(px, p, save_dir / f'{prefix}P_curve.png', names, ylabel='Precision') + plot_mc_curve(px, r, save_dir / f'{prefix}R_curve.png', names, ylabel='Recall') i = smooth(f1.mean(0), 0.1).argmax() # max F1 index p, r, f1 = p[:, i], r[:, i], f1[:, i] diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index dc31332..ad9b98b 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/ultralytics/yolo/v8/detect/train.py @@ -63,16 +63,15 @@ class DetectionTrainer(BaseTrainer): return dict(zip(keys, loss_items)) if loss_items is not None else keys def progress_string(self): - return ('\n' + '%11s' * 7) % \ - ('Epoch', 'GPU_mem', *self.loss_names, 'Instances', 'Size') + return ('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', *self.loss_names, 'Instances', 'Size') def plot_training_samples(self, batch, ni): - images = batch["img"] - cls = batch["cls"].squeeze(-1) - bboxes = batch["bboxes"] - paths = batch["im_file"] - batch_idx = batch["batch_idx"] - plot_images(images, batch_idx, cls, bboxes, paths=paths, fname=self.save_dir / f"train_batch{ni}.jpg") + plot_images(images=batch["img"], + batch_idx=batch["batch_idx"], + cls=batch["cls"].squeeze(-1), + bboxes=batch["bboxes"], + paths=batch["im_file"], + fname=self.save_dir / f"train_batch{ni}.jpg") def plot_metrics(self): plot_results(file=self.csv) # save results.png diff --git a/ultralytics/yolo/v8/segment/train.py b/ultralytics/yolo/v8/segment/train.py index 5e7f00b..2d9694e 100644 --- a/ultralytics/yolo/v8/segment/train.py +++ b/ultralytics/yolo/v8/segment/train.py @@ -214,8 +214,7 @@ class SegmentationTrainer(DetectionTrainer): return dict(zip(keys, loss_items)) if loss_items is not None else keys def progress_string(self): - return ('\n' + '%11s' * 8) % \ - ('Epoch', 'GPU_mem', *self.loss_names, 'Instances', 'Size') + return ('\n' + '%11s' * 8) % ('Epoch', 'GPU_mem', *self.loss_names, 'Instances', 'Size') def plot_training_samples(self, batch, ni): images = batch["img"]