Rename opset_version to opset (#4135)

This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-07-24 13:19:12 +02:00
1 parent 2e538443b7
commit 264be1a616
1 file changed
+5 -5
+5 -5
View File
@@ -38,7 +38,7 @@ def export_torchscript(model, img, file, optimize):
print(f'{prefix} export failure: {e}') 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 # ONNX model export
prefix = colorstr('ONNX:') prefix = colorstr('ONNX:')
try: 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__}...') print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
f = file.with_suffix('.onnx') 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, training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=not train, do_constant_folding=not train,
input_names=['images'], input_names=['images'],
@@ -108,7 +108,7 @@ def run(weights='./yolov5s.pt', # weights path
optimize=False, # TorchScript: optimize for mobile optimize=False, # TorchScript: optimize for mobile
dynamic=False, # ONNX: dynamic axes dynamic=False, # ONNX: dynamic axes
simplify=False, # ONNX: simplify model simplify=False, # ONNX: simplify model
opset_version=12, # ONNX: opset version opset=12, # ONNX: opset version
): ):
t = time.time() t = time.time()
include = [x.lower() for x in include] include = [x.lower() for x in include]
@@ -149,7 +149,7 @@ def run(weights='./yolov5s.pt', # weights path
if 'torchscript' in include: if 'torchscript' in include:
export_torchscript(model, img, file, optimize) export_torchscript(model, img, file, optimize)
if 'onnx' in include: 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: if 'coreml' in include:
export_coreml(model, img, file) 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('--optimize', action='store_true', help='TorchScript: optimize for mobile')
parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes') 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('--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() opt = parser.parse_args()
return opt return opt