split onnx export into torch.export then torch.onnx.export with the "torch.export"-ed program. This is a workaround for a bug introduced in new release of torch version 2.10.0

Author: me

src/test.py

@@ -250,7 +250,8 @@ def get_model(model: str, variant: str = ''):
         has_obj = False
 
     elif model == 'yolov7':
-        load_from_yolov7_official(net, '../weights/yolov7.pt')
+        if os.path.exists('../weights/yolov7.pt'):
+            load_from_yolov7_official(net, '../weights/yolov7.pt')
 
     return net, has_obj
 
@@ -272,12 +273,15 @@ def export(model: str, variant: str = '', onnx_path:str = '/tmp/model.onnx'):
     _ = net(x) # warmup all the einops kernels
 
     print(bcolors.OKGREEN, f"Exporting {type(net).__name__} ...", bcolors.ENDC)
-    torch.onnx.export(net, (x,), dynamo=True, opset_version=23,
-                      input_names=['img'],
-                      output_names=['preds'],
-                      dynamic_shapes={'x' : (Dim.DYNAMIC, Dim.STATIC, Dim.DYNAMIC, Dim.DYNAMIC)}).save(onnx_path)
+    prog = torch.export.export(net, (x,), dynamic_shapes={'x' : (Dim.DYNAMIC, Dim.STATIC, Dim.DYNAMIC, Dim.DYNAMIC)})
     print(bcolors.OKGREEN, f"Exporting {type(net).__name__} ... Done", bcolors.ENDC)
 
+    print(bcolors.OKGREEN, f"ONNX {type(net).__name__} ...", bcolors.ENDC)
+    torch.onnx.export(prog, dynamo=True, opset_version=23,
+                      input_names=['img'],
+                      output_names=['preds']).save(onnx_path)
+    print(bcolors.OKGREEN, f"ONNX {type(net).__name__} ... Done", bcolors.ENDC)
+
     print(bcolors.OKGREEN, f"Slimming {type(net).__name__} ...", bcolors.ENDC)
     model           = onnx.load(onnx_path)
     slimmed_model   = onnxslim.slim(model)
@@ -369,3 +373,8 @@ def export(model: str, variant: str = '', onnx_path:str = '/tmp/model.onnx'):
 # export('yolov12', 'm')
 # export('yolov12', 'l')
 # export('yolov12', 'x')
+# export('yolov26', 'n')
+# export('yolov26', 's')
+# export('yolov26', 'm')
+# export('yolov26', 'l')
+# export('yolov26', 'x')

src/unit.py

@@ -43,10 +43,10 @@ def test_export(net: YoloBase):
     net = net.eval()
     x = torch.randn(4, 3, 640, 640)
     _ = net(x) # compile einops kernels just in case
-    torch.onnx.export(net, (x,), dynamo=True, opset_version=23,
+    prog = torch.export.export(net, (x,), dynamic_shapes={'x' : (Dim.DYNAMIC, Dim.STATIC, Dim.DYNAMIC, Dim.DYNAMIC)})
+    torch.onnx.export(prog, dynamo=True, opset_version=23,
                       input_names=['img'],
-                      output_names=['preds'],
-                      dynamic_shapes={'x' : (Dim.DYNAMIC, Dim.STATIC, Dim.DYNAMIC, Dim.DYNAMIC)}).save('/tmp/model.onnx')
+                      output_names=['preds']).save('/tmp/model.onnx')
     netOrt  = ort.InferenceSession('/tmp/model.onnx', providers=['CPUExecutionProvider'])
     x       = torch.randn(2, 3, 576, 768)
     preds1  = net(x)