From 63a714e6e64836aa2672f39a3a9b69cdbc923ca3 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 24 Jul 2021 13:16:29 +0200
Subject: [PATCH] Rename `opset_version` to `opset`

---
 export.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/export.py b/export.py
index 1e3a5ed6a4af..ceee15e1644f 100644
--- a/export.py
+++ b/export.py
@@ -38,7 +38,7 @@ def export_torchscript(model, img, file, optimize):
         print(f'{prefix} export failure: {e}')
 
 
-def export_onnx(model, img, file, opset_version, train, dynamic, simplify):
+def export_onnx(model, img, file, opset, train, dynamic, simplify):
     # ONNX model export
     prefix = colorstr('ONNX:')
     try:
@@ -47,7 +47,7 @@ def export_onnx(model, img, file, opset_version, train, dynamic, simplify):
 
         print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
         f = file.with_suffix('.onnx')
-        torch.onnx.export(model, img, f, verbose=False, opset_version=opset_version,
+        torch.onnx.export(model, img, f, verbose=False, opset_version=opset,
                           training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                           do_constant_folding=not train,
                           input_names=['images'],
@@ -108,7 +108,7 @@ def run(weights='./yolov5s.pt',  # weights path
         optimize=False,  # TorchScript: optimize for mobile
         dynamic=False,  # ONNX: dynamic axes
         simplify=False,  # ONNX: simplify model
-        opset_version=12,  # ONNX: opset version
+        opset=12,  # ONNX: opset version
         ):
     t = time.time()
     include = [x.lower() for x in include]
@@ -149,7 +149,7 @@ def run(weights='./yolov5s.pt',  # weights path
     if 'torchscript' in include:
         export_torchscript(model, img, file, optimize)
     if 'onnx' in include:
-        export_onnx(model, img, file, opset_version, train, dynamic, simplify)
+        export_onnx(model, img, file, opset, train, dynamic, simplify)
     if 'coreml' in include:
         export_coreml(model, img, file)
 
@@ -170,7 +170,7 @@ def parse_opt():
     parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
     parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes')
     parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
-    parser.add_argument('--opset-version', type=int, default=12, help='ONNX: opset version')
+    parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
     opt = parser.parse_args()
     return opt