From befdf6ee64b5a75289ef4a1785eb6f432d1cc044 Mon Sep 17 00:00:00 2001 From: Silmeria112 Date: Fri, 29 Oct 2021 11:10:12 +0800 Subject: [PATCH] update gen_wts.py for yolov5 to support custom model with different anchor from official model(#788) --- yolov5/gen_wts.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/yolov5/gen_wts.py b/yolov5/gen_wts.py index 2b86e7c..1453465 100644 --- a/yolov5/gen_wts.py +++ b/yolov5/gen_wts.py @@ -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: