haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

modify load model for ema (#929)

Browse Source 
Co-authored-by: 王光强 <wangguangqiang@cethik.com>
This commit is contained in:
AdanWang 2022-03-11 13:44:17 +08:00 committed by GitHub
parent c2a07ff7ac
commit 70059d1802
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
yolov5/gen_wts.py
View File

@ -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]