Skip to content

Commit

Permalink
Rename opset_version to opset (ultralytics#4135)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed Jul 24, 2021
1 parent 80999a3 commit 898173e
Showing 1 changed file with 5 additions and 5 deletions.
10 changes: 5 additions & 5 deletions export.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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'],
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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)

Expand All @@ -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

Expand Down

0 comments on commit 898173e

Please sign in to comment.