Add onnxsim before measure (#3)

ENOT-AutoDL · Sep 5, 2023 · ba1d3b4 · ba1d3b4
1 parent 5958fbd
commit ba1d3b4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/measure.py b/measure.py
@@ -3,6 +3,7 @@
 import onnx
 import torch
 from enot_latency_server.client import measure_latency_remote
+from onnxsim import simplify
 
 from prune import measure_latency_on_server
 
@@ -26,6 +27,7 @@
 
     if args.onnx:
         onnx_model = onnx.load(args.onnx)
+        onnx_model, _ = simplify(onnx_model)
         latency = measure_latency_remote(onnx_model.SerializeToString(), host=args.host, port=args.port)
         print(latency)
         exit()

diff --git a/requirements.txt b/requirements.txt
@@ -3,6 +3,7 @@ enot-autodl==3.3.2
 enot-latency-server==1.2.0
 imagesize
 onnxruntime-gpu
+onnxsim
 opencv-python
 pathspec
 scikit-learn