From 70059d18023ee3ac07f743365f8ce95d6c779cd0 Mon Sep 17 00:00:00 2001 From: AdanWang <32757980+AdanWang@users.noreply.github.com> Date: Fri, 11 Mar 2022 13:44:17 +0800 Subject: [PATCH] modify load model for ema (#929) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 王光强 --- yolov5/gen_wts.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/yolov5/gen_wts.py b/yolov5/gen_wts.py index 1453465..5ae9497 100644 --- a/yolov5/gen_wts.py +++ b/yolov5/gen_wts.py @@ -27,7 +27,8 @@ pt_file, wts_file = parse_args() # Initialize device = select_device('cpu') # Load model -model = torch.load(pt_file, map_location=device)['model'].float() # load to FP32 +model = torch.load(pt_file, map_location=device) # load to FP32 +model = model['ema' if model.get('ema') else 'model'].float() # update anchor_grid info anchor_grid = model.model[-1].anchors * model.model[-1].stride[...,None,None]