Rename opset_version to opset (#4135)
This commit is contained in:
1 parent
2e538443b7
commit
264be1a616
1 file changed
+5
-5
@@ -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
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user