add weights_only=False for the latest torch

mpj1234 · mpj1234 · commit 09be781e147b · 2025-12-16T20:58:35.000+08:00
diff --git a/yolov10/gen_wts.py b/yolov10/gen_wts.py
@@ -38,9 +38,9 @@ def parse_args():
 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, weights_only=False)['model'].float()  # load to FP32
 # If the training is not finished, the model will be interrupted.
-# model = torch.load(pt_file, map_location=device)['ema'].float()  # load to FP32
+# model = torch.load(pt_file, map_location=device, weights_only=False)['ema'].float()  # load to FP32
 
 model.to(device).eval()
 
diff --git a/yolov3-spp/gen_wts.py b/yolov3-spp/gen_wts.py
@@ -1,13 +1,14 @@
 import struct
 import sys
-from models import *
-from utils.utils import *
+import torch
+from models import *  # noqa: F403
+from utils.utils import *  # noqa: F403
 
-model = Darknet('cfg/yolov3-spp.cfg', (416, 416))
+model = Darknet('cfg/yolov3-spp.cfg', (416, 416))  # noqa: F405
 weights = sys.argv[1]
 dev = '0'
-device = torch_utils.select_device(dev)
-model.load_state_dict(torch.load(weights, map_location=device)['model'])
+device = torch_utils.select_device(dev)  # noqa: F405
+model.load_state_dict(torch.load(weights, map_location=device, weights_only=False)['model'])
 
 
 with open('yolov3-spp_ultralytics68.wts', 'w') as f:
@@ -17,6 +18,5 @@
         f.write('{} {} '.format(k, len(vr)))
         for vv in vr:
             f.write(' ')
-            f.write(struct.pack('>f',float(vv)).hex())
+            f.write(struct.pack('>f', float(vv)).hex())
         f.write('\n')
-
diff --git a/yolov3-tiny/gen_wts.py b/yolov3-tiny/gen_wts.py
@@ -1,15 +1,16 @@
 import struct
+import torch
 import sys
-from models import *
-from utils.utils import *
+from models import *  # noqa: F403
+from utils.utils import *  # noqa: F403
 
-model = Darknet('cfg/yolov3-tiny.cfg', (608, 608))
+model = Darknet('cfg/yolov3-tiny.cfg', (608, 608))  # noqa: F405
 weights = sys.argv[1]
-device = torch_utils.select_device('0')
+device = torch_utils.select_device('0')  # noqa: F405
 if weights.endswith('.pt'):  # pytorch format
-    model.load_state_dict(torch.load(weights, map_location=device)['model'])
+    model.load_state_dict(torch.load(weights, map_location=device, weights_only=False)['model'])
 else:  # darknet format
-    load_darknet_weights(model, weights)
+    load_darknet_weights(model, weights)  # noqa: F405
 model = model.eval()
 
 with open('yolov3-tiny.wts', 'w') as f:
@@ -19,6 +20,5 @@
         f.write('{} {} '.format(k, len(vr)))
         for vv in vr:
             f.write(' ')
-            f.write(struct.pack('>f',float(vv)).hex())
+            f.write(struct.pack('>f', float(vv)).hex())
         f.write('\n')
-
diff --git a/yolov3/gen_wts.py b/yolov3/gen_wts.py
@@ -1,15 +1,16 @@
 import struct
 import sys
-from models import *
-from utils.utils import *
+import torch
+from models import *  # noqa: F403
+from utils.utils import *  # noqa: F403
 
-model = Darknet('cfg/yolov3.cfg', (608, 608))
+model = Darknet('cfg/yolov3.cfg', (608, 608))  # noqa: F405
 weights = sys.argv[1]
-device = torch_utils.select_device('0')
+device = torch_utils.select_device('0')  # noqa: F405
 if weights.endswith('.pt'):  # pytorch format
-    model.load_state_dict(torch.load(weights, map_location=device)['model'])
+    model.load_state_dict(torch.load(weights, map_location=device, weights_only=False)['model'])
 else:  # darknet format
-    load_darknet_weights(model, weights)
+    load_darknet_weights(model, weights)  # noqa: F405
 model = model.eval()
 
 with open('yolov3.wts', 'w') as f:
@@ -19,6 +20,5 @@
         f.write('{} {} '.format(k, len(vr)))
         for vv in vr:
             f.write(' ')
-            f.write(struct.pack('>f',float(vv)).hex())
+            f.write(struct.pack('>f', float(vv)).hex())
         f.write('\n')
-
diff --git a/yolov4/gen_wts.py b/yolov4/gen_wts.py
@@ -1,15 +1,16 @@
 import struct
 import sys
-from models import *
-from utils.utils import *
+import torch
+from models import *  # noqa: F403
+from utils.utils import *  # noqa: F403
 
-model = Darknet('cfg/yolov4.cfg', (608, 608))
+model = Darknet('cfg/yolov4.cfg', (608, 608))  # noqa: F405
 weights = sys.argv[1]
-device = torch_utils.select_device('0')
+device = torch_utils.select_device('0')  # noqa: F405
 if weights.endswith('.pt'):  # pytorch format
-    model.load_state_dict(torch.load(weights, map_location=device)['model'])
+    model.load_state_dict(torch.load(weights, map_location=device, weights_only=False)['model'])
 else:  # darknet format
-    load_darknet_weights(model, weights)
+    load_darknet_weights(model, weights)  # noqa: F405
 
 with open('yolov4.wts', 'w') as f:
     f.write('{}\n'.format(len(model.state_dict().keys())))
@@ -18,6 +19,5 @@
         f.write('{} {} '.format(k, len(vr)))
         for vv in vr:
             f.write(' ')
-            f.write(struct.pack('>f',float(vv)).hex())
+            f.write(struct.pack('>f', float(vv)).hex())
         f.write('\n')
-
diff --git a/yolov5/gen_wts.py b/yolov5/gen_wts.py
@@ -1,4 +1,3 @@
-import sys
 import argparse
 import os
 import struct
@@ -33,7 +32,7 @@ def parse_args():
 # Load model
 print(f'Loading {pt_file}')
 device = select_device('cpu')
-model = torch.load(pt_file, map_location=device)  # Load FP32 weights
+model = torch.load(pt_file, map_location=device, weights_only=False)  # Load FP32 weights
 model = model['ema' if model.get('ema') else 'model'].float()
 
 if m_type in ['detect', 'seg']:
diff --git a/yolov7/gen_wts.py b/yolov7/gen_wts.py
@@ -1,4 +1,3 @@
-import sys
 import argparse
 import os
 import struct
@@ -27,13 +26,14 @@ def 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, weights_only=False)['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", anchor_grid)  # The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight.
+# 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)
 
 model.to(device).eval()
 
diff --git a/yolov8/gen_wts.py b/yolov8/gen_wts.py
@@ -37,7 +37,7 @@ def parse_args():
 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, weights_only=False)['model'].float()  # load to FP32
 
 if m_type in ['detect', 'seg', 'pose']:
     anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]
diff --git a/yolov8/yolov8_trt10/gen_wts.py b/yolov8/yolov8_trt10/gen_wts.py
@@ -37,7 +37,7 @@ def parse_args():
 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, weights_only=False)['model'].float()  # load to FP32
 
 if m_type in ['detect', 'seg', 'pose']:
     anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]
diff --git a/yolov9/gen_wts.py b/yolov9/gen_wts.py
@@ -1,10 +1,10 @@
-import sys
 import argparse
 import os
 import struct
 import torch
 from utils.torch_utils import select_device
 
+
 def parse_args():
     parser = argparse.ArgumentParser(description='Convert .pt file to .wts')
     parser.add_argument('-w', '--weights', default='yolov9-e.pt',
@@ -25,13 +25,14 @@ def parse_args():
             os.path.splitext(os.path.basename(args.weights))[0] + '.wts')
     return args.weights, args.output, args.type
 
+
 pt_file, wts_file, m_type = parse_args()
 print(f'Generating .wts for {m_type} model')
 
 # Load model
 print(f'Loading {pt_file}')
 device = select_device('cpu')
-model = torch.load(pt_file, map_location=device)  # Load FP32 weights
+model = torch.load(pt_file, map_location=device, weights_only=False)  # Load FP32 weights
 model = model["model"].float()
 
 if m_type in ['detect', 'seg']:
