fix genwts file close

Author: wang-xinyu

yolov3-spp/gen_wts.py

@@ -10,13 +10,13 @@
 model.load_state_dict(torch.load(weights, map_location=device)['model'])
 
 
-f = open('yolov3-spp_ultralytics68.wts', 'w')
-f.write('{}\n'.format(len(model.state_dict().keys())))
-for k, v in model.state_dict().items():
-    vr = v.reshape(-1).cpu().numpy()
-    f.write('{} {} '.format(k, len(vr)))
-    for vv in vr:
-        f.write(' ')
-        f.write(struct.pack('>f',float(vv)).hex())
-    f.write('\n')
+with open('yolov3-spp_ultralytics68.wts', 'w') as f:
+    f.write('{}\n'.format(len(model.state_dict().keys())))
+    for k, v in model.state_dict().items():
+        vr = v.reshape(-1).cpu().numpy()
+        f.write('{} {} '.format(k, len(vr)))
+        for vv in vr:
+            f.write(' ')
+            f.write(struct.pack('>f',float(vv)).hex())
+        f.write('\n')
 

yolov3-tiny/gen_wts.py

@@ -12,13 +12,13 @@
     load_darknet_weights(model, weights)
 model = model.eval()
 
-f = open('yolov3-tiny.wts', 'w')
-f.write('{}\n'.format(len(model.state_dict().keys())))
-for k, v in model.state_dict().items():
-    vr = v.reshape(-1).cpu().numpy()
-    f.write('{} {} '.format(k, len(vr)))
-    for vv in vr:
-        f.write(' ')
-        f.write(struct.pack('>f',float(vv)).hex())
-    f.write('\n')
+with open('yolov3-tiny.wts', 'w') as f:
+    f.write('{}\n'.format(len(model.state_dict().keys())))
+    for k, v in model.state_dict().items():
+        vr = v.reshape(-1).cpu().numpy()
+        f.write('{} {} '.format(k, len(vr)))
+        for vv in vr:
+            f.write(' ')
+            f.write(struct.pack('>f',float(vv)).hex())
+        f.write('\n')
 

yolov3/gen_wts.py

@@ -12,13 +12,13 @@
     load_darknet_weights(model, weights)
 model = model.eval()
 
-f = open('yolov3.wts', 'w')
-f.write('{}\n'.format(len(model.state_dict().keys())))
-for k, v in model.state_dict().items():
-    vr = v.reshape(-1).cpu().numpy()
-    f.write('{} {} '.format(k, len(vr)))
-    for vv in vr:
-        f.write(' ')
-        f.write(struct.pack('>f',float(vv)).hex())
-    f.write('\n')
+with open('yolov3.wts', 'w') as f:
+    f.write('{}\n'.format(len(model.state_dict().keys())))
+    for k, v in model.state_dict().items():
+        vr = v.reshape(-1).cpu().numpy()
+        f.write('{} {} '.format(k, len(vr)))
+        for vv in vr:
+            f.write(' ')
+            f.write(struct.pack('>f',float(vv)).hex())
+        f.write('\n')
 

yolov4/gen_wts.py

@@ -11,13 +11,13 @@
 else:  # darknet format
     load_darknet_weights(model, weights)
 
-f = open('yolov4.wts', 'w')
-f.write('{}\n'.format(len(model.state_dict().keys())))
-for k, v in model.state_dict().items():
-    vr = v.reshape(-1).cpu().numpy()
-    f.write('{} {} '.format(k, len(vr)))
-    for vv in vr:
-        f.write(' ')
-        f.write(struct.pack('>f',float(vv)).hex())
-    f.write('\n')
+with open('yolov4.wts', 'w') as f:
+    f.write('{}\n'.format(len(model.state_dict().keys())))
+    for k, v in model.state_dict().items():
+        vr = v.reshape(-1).cpu().numpy()
+        f.write('{} {} '.format(k, len(vr)))
+        for vv in vr:
+            f.write(' ')
+            f.write(struct.pack('>f',float(vv)).hex())
+        f.write('\n')
 

yolov5/gen_wts.py

@@ -8,12 +8,12 @@
 model = torch.load('weights/yolov5s.pt', map_location=device)['model'].float()  # load to FP32
 model.to(device).eval()
 
-f = open('yolov5s.wts', 'w')
-f.write('{}\n'.format(len(model.state_dict().keys())))
-for k, v in model.state_dict().items():
-    vr = v.reshape(-1).cpu().numpy()
-    f.write('{} {} '.format(k, len(vr)))
-    for vv in vr:
-        f.write(' ')
-        f.write(struct.pack('>f',float(vv)).hex())
-    f.write('\n')
+with open('yolov5s.wts', 'w') as f:
+    f.write('{}\n'.format(len(model.state_dict().keys())))
+    for k, v in model.state_dict().items():
+        vr = v.reshape(-1).cpu().numpy()
+        f.write('{} {} '.format(k, len(vr)))
+        for vv in vr:
+            f.write(' ')
+            f.write(struct.pack('>f',float(vv)).hex())
+        f.write('\n')