From e5928b095a8bc93e1416d6b1b3f51183995ab3df Mon Sep 17 00:00:00 2001 From: AlexeyAB84 Date: Fri, 22 Jul 2022 17:04:38 +0300 Subject: [PATCH] minor fix --- .idea/.gitignore | 3 -- .idea/inspectionProfiles/Project_Default.xml | 46 ------------------- .../inspectionProfiles/profiles_settings.xml | 6 --- .idea/misc.xml | 4 -- .idea/modules.xml | 8 ---- .idea/vcs.xml | 6 --- .idea/yolov7.iml | 12 ----- README.md | 13 +++++- export.py | 34 +++++++------- 9 files changed, 29 insertions(+), 103 deletions(-) delete mode 100644 .idea/.gitignore delete mode 100644 .idea/inspectionProfiles/Project_Default.xml delete mode 100644 .idea/inspectionProfiles/profiles_settings.xml delete mode 100644 .idea/misc.xml delete mode 100644 .idea/modules.xml delete mode 100644 .idea/vcs.xml delete mode 100644 .idea/yolov7.iml diff --git a/.idea/.gitignore b/.idea/.gitignore deleted file mode 100644 index 359bb53..0000000 --- a/.idea/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -# 默认忽略的文件 -/shelf/ -/workspace.xml diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml deleted file mode 100644 index 28c863b..0000000 --- a/.idea/inspectionProfiles/Project_Default.xml +++ /dev/null @@ -1,46 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml deleted file mode 100644 index 105ce2d..0000000 --- a/.idea/inspectionProfiles/profiles_settings.xml +++ /dev/null @@ -1,6 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml deleted file mode 100644 index d1e22ec..0000000 --- a/.idea/misc.xml +++ /dev/null @@ -1,4 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml deleted file mode 100644 index ecffae6..0000000 --- a/.idea/modules.xml +++ /dev/null @@ -1,8 +0,0 @@ - - - - - - - - \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml deleted file mode 100644 index 94a25f7..0000000 --- a/.idea/vcs.xml +++ /dev/null @@ -1,6 +0,0 @@ - - - - - - \ No newline at end of file diff --git a/.idea/yolov7.iml b/.idea/yolov7.iml deleted file mode 100644 index 8b8c395..0000000 --- a/.idea/yolov7.iml +++ /dev/null @@ -1,12 +0,0 @@ - - - - - - - - - - \ No newline at end of file diff --git a/README.md b/README.md index 1892b1b..38113bf 100644 --- a/README.md +++ b/README.md @@ -153,9 +153,18 @@ python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source inferen ## Export -Use the args `--include-nms` can to export end to end onnx model which include the `EfficientNMS`. +Tested with: Python 3.7.13 and Pytorch 1.12.0+cu113 +Pytorch to ONNX, use `--include-nms` flag for the end-to-end ONNX model with `EfficientNMS`. ```shell -python models/export.py --weights yolov7.pt --grid --include-nms +wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt +python export.py --weights yolov7-tiny.pt --grid --include-nms +``` + +ONNX to TensorRT +```shell +git clone https://github.com/Linaom1214/tensorrt-python.git +cd tensorrt-python +python export.py -o yolov7-tiny.onnx -e yolov7-tiny-nms.trt -p fp16 ``` ## Citation diff --git a/export.py b/export.py index fd4975f..f9a6a89 100644 --- a/export.py +++ b/export.py @@ -79,23 +79,17 @@ if __name__ == '__main__': dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640) 'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None) - if opt.include_nms: - print('Registering NMS plugin...') - mo = RegisterNMS(f) - mo.register_nms() - mo.save(f) - else: - # Checks - onnx_model = onnx.load(f) # load onnx model - onnx.checker.check_model(onnx_model) # check onnx model - # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model + # Checks + onnx_model = onnx.load(f) # load onnx model + onnx.checker.check_model(onnx_model) # check onnx model + # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model - # # Metadata - # d = {'stride': int(max(model.stride))} - # for k, v in d.items(): - # meta = onnx_model.metadata_props.add() - # meta.key, meta.value = k, str(v) - # onnx.save(onnx_model, f) + # # Metadata + # d = {'stride': int(max(model.stride))} + # for k, v in d.items(): + # meta = onnx_model.metadata_props.add() + # meta.key, meta.value = k, str(v) + # onnx.save(onnx_model, f) if opt.simplify: try: @@ -104,9 +98,17 @@ if __name__ == '__main__': print('\nStarting to simplify ONNX...') onnx_model, check = onnxsim.simplify(onnx_model) assert check, 'assert check failed' + onnx.save(onnx_model, f) except Exception as e: print(f'Simplifier failure: {e}') print('ONNX export success, saved as %s' % f) + + if opt.include_nms: + print('Registering NMS plugin for ONNX...') + mo = RegisterNMS(f) + mo.register_nms() + mo.save(f) + except Exception as e: print('ONNX export failure: %s' % e) # CoreML export