add onnx-sim handle

requirements.txt

@@ -1,4 +1,8 @@
-numpy==1.19.5
+numpy==1.23.0
+onnx==1.12.0
+onnx_simplifier==0.3.10
+onnxruntime==1.11.1
+onnxsim==0.4.0
 opencv_python==4.2.0.34
 pycocotools==2.0.4
 PyYAML==6.0

test.py

@@ -1,7 +1,9 @@
 import os
 import cv2
+import onnx
 import time
 import argparse
+from onnxsim import simplify
 
 import torch
 from utils.tool import *
@@ -55,12 +57,18 @@ if __name__ == '__main__':
 
     # 导出onnx模型
     if opt.onnx:
-        torch.onnx.export(model,                    #model being run
-                          img,                 # model input (or a tuple for multiple inputs)
-                          "./FastestDet.onnx",             # where to save the model (can be a file or file-like object)
+        torch.onnx.export(model,                     # model being run
+                          img,                       # model input (or a tuple for multiple inputs)
+                          "./FastestDet.onnx",       # where to save the model (can be a file or file-like object)
                           export_params=True,        # store the trained parameter weights inside the model file
                           opset_version=11,          # the ONNX version to export the model to
                           do_constant_folding=True)  # whether to execute constant folding for optimization
+        # onnx-sim
+        onnx_model = onnx.load("./FastestDet.onnx")  # load onnx model
+        model_simp, check = simplify(onnx_model)
+        assert check, "Simplified ONNX model could not be validated"
+        print("onnx sim sucess...")
+        onnx.save(model_simp, "./FastestDet.onnx")                  
 
     # 导出torchscript模型
     if opt.torchscript: