diff --git a/segment/train.py b/segment/train.py index 5121c5fa784a..26f0d0c13c78 100644 --- a/segment/train.py +++ b/segment/train.py @@ -651,7 +651,7 @@ def main(opt, callbacks=Callbacks()): results = train(hyp.copy(), opt, device, callbacks) callbacks = Callbacks() # Write mutation results - print_mutation(results, hyp.copy(), save_dir, opt.bucket) + print_mutation(KEYS, results, hyp.copy(), save_dir, opt.bucket) # Plot results plot_evolve(evolve_csv) diff --git a/train.py b/train.py index 9efece250581..177e081c8c37 100644 --- a/train.py +++ b/train.py @@ -607,7 +607,9 @@ def main(opt, callbacks=Callbacks()): results = train(hyp.copy(), opt, device, callbacks) callbacks = Callbacks() # Write mutation results - print_mutation(results, hyp.copy(), save_dir, opt.bucket) + keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss', + 'val/obj_loss', 'val/cls_loss') + print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket) # Plot results plot_evolve(evolve_csv) diff --git a/utils/general.py b/utils/general.py index d31b043a113e..e2faca9dbf2a 100644 --- a/utils/general.py +++ b/utils/general.py @@ -957,11 +957,10 @@ def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_op LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB") -def print_mutation(results, hyp, save_dir, bucket, prefix=colorstr('evolve: ')): +def print_mutation(keys, results, hyp, save_dir, bucket, prefix=colorstr('evolve: ')): evolve_csv = save_dir / 'evolve.csv' evolve_yaml = save_dir / 'hyp_evolve.yaml' - keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss', - 'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys()) # [results + hyps] + keys = tuple(keys) + tuple(hyp.keys()) # [results + hyps] keys = tuple(x.strip() for x in keys) vals = results + tuple(hyp.values()) n = len(keys)