From 32493ab5a5fe7a402a3dcc22b32f90e3ddb395bc Mon Sep 17 00:00:00 2001 From: SecretStar112 Date: Thu, 7 Jan 2021 08:27:22 +0800 Subject: [PATCH] W&B ID reset on training completion (#1852) * Update train.py Fix the bug of always the same W&B ID and continue overwrite with the old logging. BUG report https://github.com/ultralytics/yolov5/issues/1851 * Fix the bug of duplicate W&B ID fix the bug of https://github.com/ultralytics/yolov5/issues/1851 If we had trained on yolov5s.pt, the program will generate a new unique W&B ID. If we hadn't, the program will keep the old code, we can still use --resume aug. * Update general.py * revert train.py changes Co-authored-by: Glenn Jocher --- utils/general.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/general.py b/utils/general.py index 797587b..90265c6 100755 --- a/utils/general.py +++ b/utils/general.py @@ -361,8 +361,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non def strip_optimizer(f='weights/best.pt', s=''): # from utils.general import *; strip_optimizer() # Strip optimizer from 'f' to finalize training, optionally save as 's' x = torch.load(f, map_location=torch.device('cpu')) - x['optimizer'] = None - x['training_results'] = None + for key in 'optimizer', 'training_results', 'wandb_id': + x[key] = None x['epoch'] = -1 x['model'].half() # to FP16 for p in x['model'].parameters():