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

update gen_wts.py for yolov5 to support custom model with different anchor from official model(#788)

Browse Source
This commit is contained in:
Silmeria112 2021-10-29 11:10:12 +08:00 committed by GitHub
parent 4697890416
commit befdf6ee64
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
yolov5/gen_wts.py
View File

@ -28,8 +28,13 @@ pt_file, wts_file = parse_args()
device = select_device('cpu')
# Load model
model = torch.load(pt_file, map_location=device)['model'].float() # load to FP32
# update anchor_grid info
anchor_grid = model.model[-1].anchors * model.model[-1].stride[...,None,None]
# model.model[-1].anchor_grid = anchor_grid
delattr(model.model[-1], 'anchor_grid') # model.model[-1] is detect layer
model.model[-1].register_buffer("anchor_grid",torch.Tensor(model.yaml['anchors'])) #The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight.
model.model[-1].register_buffer("anchor_grid",anchor_grid) #The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight.
model.to(device).eval()
with open(wts_file, 'w') as f: