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

Author: Silmeria112

yolov5/gen_wts.py

@@ -28,8 +28,13 @@ def 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: