简化onnx

Author: lipengbo

Models/Backend/onnx.py

@@ -13,15 +13,14 @@ def __init__(self):
         pass
 
     @staticmethod
-    def convert(model, imgs, weights, dynamic, simplify):
+    def convert(model, imgs, weights, dynamic):
         """
         torch模型转为onnx模型
 
         model: torch模型
         imgs: [B,C,H,W]Tensor
         weights: onnx权重保存路径
         dynamic: batch轴是否设为动态维度
-        simplify: 是否简化onnx
         """
         torch.onnx.export(
             model,
@@ -37,17 +36,18 @@ def convert(model, imgs, weights, dynamic, simplify):
         )
         model_onnx = onnx.load(weights)  # load onnx model
         onnx.checker.check_model(model_onnx)  # check onnx model
-        if simplify:
-            try:
-                model_onnx, check = onnxsim.simplify(
-                    model_onnx,
-                    dynamic_input_shape=dynamic,
-                    input_shapes={"input": list(imgs.shape)} if dynamic else None,
-                )
-                assert check, "assert check failed"
-                onnx.save(model_onnx, weights)
-            except Exception as e:
-                print(f"simplifer failure: {e}")
+
+        try:
+            # 简化onnx
+            model_onnx, check = onnxsim.simplify(
+                model_onnx,
+                dynamic_input_shape=dynamic,
+                input_shapes={"input": list(imgs.shape)} if dynamic else None,
+            )
+            assert check, "assert check failed"
+            onnx.save(model_onnx, weights)
+        except Exception as e:
+            print(f"simplifer failure: {e}")
 
         print("*" * 28)
         print("ONNX export success, saved as %s" % weights)

export.py

@@ -16,7 +16,6 @@
 
     # onnx
     parser.add_argument("--torch2onnx", action="store_true", help="(可选)转为onnx")
-    parser.add_argument("--simplify", action="store_true", help="(可选)简化onnx")
     parser.add_argument("--dynamic", action="store_true", help="(可选)batch轴设为动态")
 
     # tensorrt
@@ -67,8 +66,7 @@
             model=model,
             imgs=imgs,
             weights=onnx_weights,
-            dynamic=cfg.dynamic,
-            simplify=cfg.simplify,
+            dynamic=cfg.dynamic
         )
         output_onnx = OnnxBackend.infer(weights=onnx_weights, imgs=imgs.numpy())
 
@@ -108,8 +106,10 @@
 
         from Models.Backend.mnn import MNNBackbend
 
-        MNNBackbend.convert(onnx_weights,mnn_weights,fp16=cfg.mnn_fp16)
-        output_mnn = MNNBackbend.infer(mnn_weights, imgs.numpy(),output_shape=output_onnx.shape)
+        MNNBackbend.convert(onnx_weights, mnn_weights, fp16=cfg.mnn_fp16)
+        output_mnn = MNNBackbend.infer(
+            mnn_weights, imgs.numpy(), output_shape=output_onnx.shape
+        )
     # ==========================验证结果===============================
     print("\n", "*" * 28)
     if cfg.torch2script: