diff --git a/train.py b/train.py index d5191d56017..faff65371d9 100644 --- a/train.py +++ b/train.py @@ -476,7 +476,7 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio total_training_hours = (time.time() - t0) / 3600 if RANK in {-1, 0}: - LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {total_training_hours:.3f} hours.') + LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {total_training_hours:.3f} hours.") for f in last, best: if f.exists(): strip_optimizer(f) # strip optimizers @@ -500,7 +500,7 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio if is_coco: callbacks.run("on_fit_epoch_end", list(mloss) + list(results) + lr, epoch, best_fitness, fi) - callbacks.run('on_train_end', last, best, epoch, results, total_training_hours) + callbacks.run("on_train_end", last, best, epoch, results, total_training_hours) torch.cuda.empty_cache() return results diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py index 97867f6a1f7..f1d051df254 100644 --- a/utils/loggers/__init__.py +++ b/utils/loggers/__init__.py @@ -103,7 +103,7 @@ def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, self.csv = True # always log to csv self.ndjson_console = "ndjson_console" in self.include # log ndjson to console self.ndjson_file = "ndjson_file" in self.include # log ndjson to file - self.metrics = True # log to metrics file at the end of the run + self.metrics = True # log to metrics file at the end of the run # Messages if not comet_ml: prefix = colorstr("Comet: ") @@ -304,10 +304,16 @@ def on_train_end(self, last, best, epoch, results, total_train_time): self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats="HWC") if self.metrics: - file = self.save_dir / 'metrics.json' - m={"precision":results[0],"recall":results[1],"mAP_0_5":results[2],"mAP_0_5-0_95":results[3], "total_train_hours":total_train_time} - - with open(file, 'w') as f: + file = self.save_dir / "metrics.json" + m = { + "precision": results[0], + "recall": results[1], + "mAP_0_5": results[2], + "mAP_0_5-0_95": results[3], + "total_train_hours": total_train_time, + } + + with open(file, "w") as f: f.write(json.dumps(m)) if self.wandb: