|
24 | 24 | parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/
|
25 | 25 | parser.add_argument('--img_size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
|
26 | 26 | parser.add_argument('--batch_size', type=int, default=1, help='batch size')
|
27 |
| - parser.add_argument('--simplify', action='store_true', default=False, help='simplify onnx') |
28 | 27 | parser.add_argument('--dynamic', action='store_true', default=False, help='enable dynamic axis in onnx model')
|
29 | 28 | parser.add_argument('--onnx2pb', action='store_true', default=False, help='export onnx to pb')
|
30 | 29 | parser.add_argument('--onnx_infer', action='store_true', default=True, help='onnx infer test')
|
|
88 | 87 | # Checks
|
89 | 88 | onnx_model = onnx.load(f) # load onnx model
|
90 | 89 | onnx.checker.check_model(onnx_model) # check onnx model
|
91 |
| - |
92 |
| - # https://github.com/daquexian/onnx-simplifier |
93 |
| - if opt.simplify: |
94 |
| - try: |
95 |
| - import onnxsim |
96 |
| - print(f'simplifying with onnx-simplifier {onnxsim.__version__}...') |
97 |
| - onnx_model, check = onnxsim.simplify(onnx_model, |
98 |
| - dynamic_input_shape=opt.dynamic, |
99 |
| - input_shapes={'input': list(img.shape)} if opt.dynamic else None) |
100 |
| - assert check, "simplify check failed " |
101 |
| - onnx.save(onnx_model, f) |
102 |
| - except Exception as e: |
103 |
| - print(f"simplifer failure: {e}") |
104 |
| - |
105 | 90 | print('ONNX export success, saved as %s' % f)
|
106 | 91 | # Finish
|
107 | 92 | print('\nExport complete (%.2fs). Visualize with https://github.com/lutzroeder/netron.' % (time.time() - t))
|
|
0 commit comments