From 2ce38a059254664cbb9d09d0385396cdde7bff7b Mon Sep 17 00:00:00 2001 From: Tim Date: Mon, 19 Apr 2021 01:07:39 +0200 Subject: [PATCH 1/3] Ensure dynamic export works succesfully, onnx simplifier optional --- models/export.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/models/export.py b/models/export.py index bec9194319c1..f4b876d4b6dd 100644 --- a/models/export.py +++ b/models/export.py @@ -27,6 +27,7 @@ parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') parser.add_argument('--grid', action='store_true', help='export Detect() layer grid') parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') + parser.add_argument('--onnx-simplify', action='store_true', help='simplify the exported onnx model using onnxsim') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand print(opt) @@ -87,16 +88,20 @@ # print(onnx.helper.printable_graph(model_onnx.graph)) # print # Simplify - try: - check_requirements(['onnx-simplifier']) - import onnxsim - - print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') - model_onnx, check = onnxsim.simplify(model_onnx) - assert check, 'assert check failed' - onnx.save(model_onnx, f) - except Exception as e: - print(f'{prefix} simplifier failure: {e}') + if opt.onnx_simplify: + try: + check_requirements(['onnx-simplifier']) + import onnxsim + + print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') + model_onnx, check = onnxsim.simplify( + model_onnx, + dynamic_input_shape=opt.dynamic, + input_shapes={'images': list(img.shape)} if opt.dynamic else None) + assert check, 'assert check failed' + onnx.save(model_onnx, f) + except Exception as e: + print(f'{prefix} simplifier failure: {e}') print(f'{prefix} export success, saved as {f}') except Exception as e: print(f'{prefix} export failure: {e}') From ff03e37f08a85086e1d4904ba28e129fea33b44e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 20 Apr 2021 13:38:46 +0200 Subject: [PATCH 2/3] Update export.py --- models/export.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/models/export.py b/models/export.py index f4b876d4b6dd..7cdb1a4c1963 100644 --- a/models/export.py +++ b/models/export.py @@ -21,13 +21,13 @@ if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/ + parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width parser.add_argument('--batch-size', type=int, default=1, help='batch size') - parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') parser.add_argument('--grid', action='store_true', help='export Detect() layer grid') parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--onnx-simplify', action='store_true', help='simplify the exported onnx model using onnxsim') + parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only + parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand print(opt) @@ -88,16 +88,15 @@ # print(onnx.helper.printable_graph(model_onnx.graph)) # print # Simplify - if opt.onnx_simplify: + if opt.simplify: try: check_requirements(['onnx-simplifier']) import onnxsim print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') - model_onnx, check = onnxsim.simplify( - model_onnx, - dynamic_input_shape=opt.dynamic, - input_shapes={'images': list(img.shape)} if opt.dynamic else None) + model_onnx, check = onnxsim.simplify(model_onnx, + dynamic_input_shape=opt.dynamic, + input_shapes={'images': list(img.shape)} if opt.dynamic else None) assert check, 'assert check failed' onnx.save(model_onnx, f) except Exception as e: From 1190aa0df2d3906f5b345c64bcae48d81802723a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 20 Apr 2021 13:42:52 +0200 Subject: [PATCH 3/3] add dashes --- models/export.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/models/export.py b/models/export.py index 7cdb1a4c1963..c527a47951cb 100644 --- a/models/export.py +++ b/models/export.py @@ -59,7 +59,7 @@ model.model[-1].export = not opt.grid # set Detect() layer grid export y = model(img) # dry run - # TorchScript export + # TorchScript export ----------------------------------------------------------------------------------------------- prefix = colorstr('TorchScript:') try: print(f'\n{prefix} starting export with torch {torch.__version__}...') @@ -70,7 +70,7 @@ except Exception as e: print(f'{prefix} export failure: {e}') - # ONNX export + # ONNX export ------------------------------------------------------------------------------------------------------ prefix = colorstr('ONNX:') try: import onnx @@ -105,7 +105,7 @@ except Exception as e: print(f'{prefix} export failure: {e}') - # CoreML export + # CoreML export ---------------------------------------------------------------------------------------------------- prefix = colorstr('CoreML:') try: import coremltools as ct