Unverified 提交 264be1a6 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

Rename `opset_version` to `opset` (#4135)

上级 2e538443
...@@ -38,7 +38,7 @@ def export_torchscript(model, img, file, optimize): ...@@ -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): ...@@ -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 ...@@ -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 ...@@ -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(): ...@@ -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
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论