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}')
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