From 1c9eb445589cbee999c2b66d3d18bbc564057a86 Mon Sep 17 00:00:00 2001
From: Hui Zhang <zhtclz@foxmail.com>
Date: Wed, 15 Jun 2022 03:55:24 +0000
Subject: [PATCH] onnx opt and check result

---
 .../examples/ds2_ol/onnx/local/infer_check.py  |  8 +++++++-
 .../examples/ds2_ol/onnx/local/onnx_clone.sh   |  0
 .../ds2_ol/onnx/local/onnx_infer_shape.py      |  0
 speechx/examples/ds2_ol/onnx/local/onnx_opt.sh | 18 +++++++++++++++++-
 .../ds2_ol/onnx/local/onnx_prune_model.py      |  0
 speechx/examples/ds2_ol/onnx/run.sh            |  9 +++++++++
 6 files changed, 33 insertions(+), 2 deletions(-)
 mode change 100644 => 100755 speechx/examples/ds2_ol/onnx/local/onnx_clone.sh
 mode change 100644 => 100755 speechx/examples/ds2_ol/onnx/local/onnx_infer_shape.py
 mode change 100644 => 100755 speechx/examples/ds2_ol/onnx/local/onnx_prune_model.py

diff --git a/speechx/examples/ds2_ol/onnx/local/infer_check.py b/speechx/examples/ds2_ol/onnx/local/infer_check.py
index 4debf4d33..5484b3141 100755
--- a/speechx/examples/ds2_ol/onnx/local/infer_check.py
+++ b/speechx/examples/ds2_ol/onnx/local/infer_check.py
@@ -34,6 +34,12 @@ def parse_args():
         default=".",
         help="paddle model dir."
     )
+    parser.add_argument(
+        '--model_prefix',
+        type=str,
+        default="avg_1.jit",
+        help="paddle model prefix."
+    )
     parser.add_argument(
         '--onnx_model',
         type=str,
@@ -58,7 +64,7 @@ if __name__ == '__main__':
     chunk_state_c_box = iodict['chunk_state_c_bos']
 
     # paddle
-    model = paddle.jit.load(os.path.join(FLAGS.model_dir, "avg_1.jit"))
+    model = paddle.jit.load(os.path.join(FLAGS.model_dir, FLAGS.model_prefix))
     res_chunk, res_lens, chunk_state_h, chunk_state_c = model(
         paddle.to_tensor(audio_chunk),
         paddle.to_tensor(audio_chunk_lens),
diff --git a/speechx/examples/ds2_ol/onnx/local/onnx_clone.sh b/speechx/examples/ds2_ol/onnx/local/onnx_clone.sh
old mode 100644
new mode 100755
diff --git a/speechx/examples/ds2_ol/onnx/local/onnx_infer_shape.py b/speechx/examples/ds2_ol/onnx/local/onnx_infer_shape.py
old mode 100644
new mode 100755
diff --git a/speechx/examples/ds2_ol/onnx/local/onnx_opt.sh b/speechx/examples/ds2_ol/onnx/local/onnx_opt.sh
index 470a9a317..db039a6ef 100755
--- a/speechx/examples/ds2_ol/onnx/local/onnx_opt.sh
+++ b/speechx/examples/ds2_ol/onnx/local/onnx_opt.sh
@@ -1,4 +1,20 @@
 #!/bin/bash
 
+set -e
+
+if [ $# != 3 ];then
+    # ./local/onnx_opt.sh model.old.onnx model.opt.onnx  "audio_chunk:1,-1,161  audio_chunk_lens:1 chunk_state_c_box:5,1,1024 chunk_state_h_box:5,1,1024"                                                 
+    echo "usage: $0 onnx.model.in onnx.model.out input_shape "
+    exit 1
+fi
+
 # onnx optimizer
-onnx-simplifier
\ No newline at end of file
+pip install onnx-simplifier
+
+in=$1
+out=$2
+input_shape=$3
+
+check_n=3
+
+onnxsim $in $2 $check_n --dynamic-input-shape  --input-shape $input_shape 
\ No newline at end of file
diff --git a/speechx/examples/ds2_ol/onnx/local/onnx_prune_model.py b/speechx/examples/ds2_ol/onnx/local/onnx_prune_model.py
old mode 100644
new mode 100755
diff --git a/speechx/examples/ds2_ol/onnx/run.sh b/speechx/examples/ds2_ol/onnx/run.sh
index b7c7e2fbd..9ed2aa809 100755
--- a/speechx/examples/ds2_ol/onnx/run.sh
+++ b/speechx/examples/ds2_ol/onnx/run.sh
@@ -53,6 +53,15 @@ fi
 if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ];then
     # to onnx
    ./local/tonnx.sh $dir $model $param $exp/model.onnx
+
    ./local/infer_check.py --input_file 'static_ds2online_inputs.pickle' --model_dir $dir --onnx_model $exp/model.onnx
 fi
 
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ] ;then
+    input_shape="audio_chunk:1,-1,161  audio_chunk_lens:1 chunk_state_c_box:5,1,1024 chunk_state_h_box:5,1,1024"   
+    # simplifying onnx model
+    ./local/onnx_opt.sh $exp/model.onnx $exp/model.opt.onnx  $input_shape    
+
+    ./local/infer_check.py --input_file 'static_ds2online_inputs.pickle' --model_dir $dir --onnx_model $exp/model.opt.onnx
+fi
\ No newline at end of file