reconstruct the export function and the run.sh in aishell and librispeech

deepspeech/exps/deepspeech2/model.py

@@ -367,20 +367,8 @@ class DeepSpeech2Tester(DeepSpeech2Trainer):
                                             dtype='int64'),  # audio_length, [B]
                 ])
         elif self.args.model_type == 'online':
-            static_model = paddle.jit.to_static(
-                infer_model,
-                input_spec=[
-                    paddle.static.InputSpec(
-                        shape=[None, None,
-                               feat_dim],  #[B, chunk_size, feat_dim]
-                        dtype='float32'),  # audio, [B,T,D]
-                    paddle.static.InputSpec(shape=[None],
-                                            dtype='int64'),  # audio_length, [B]
-                    paddle.static.InputSpec(
-                        shape=[None, None, None], dtype='float32'),
-                    paddle.static.InputSpec(
-                        shape=[None, None, None], dtype='float32')
-                ])
+            static_model = DeepSpeech2InferModelOnline.export(infer_model,
+                                                              feat_dim)
         else:
             raise Exception("wrong model type")
         logger.info(f"Export code: {static_model.forward.code}")

deepspeech/models/ds2_online/deepspeech2.py

@@ -424,3 +424,20 @@ class DeepSpeech2InferModelOnline(DeepSpeech2ModelOnline):
             audio_chunk, audio_chunk_lens, chunk_state_h_box, chunk_state_c_box)
         probs_chunk = self.decoder.softmax(eouts_chunk)
         return probs_chunk, eouts_chunk_lens, final_state_h_box, final_state_c_box
+
+    @classmethod
+    def export(self, infer_model, feat_dim):
+        static_model = paddle.jit.to_static(
+            infer_model,
+            input_spec=[
+                paddle.static.InputSpec(
+                    shape=[None, None, feat_dim],  #[B, chunk_size, feat_dim]
+                    dtype='float32'),  # audio, [B,T,D]
+                paddle.static.InputSpec(shape=[None],
+                                        dtype='int64'),  # audio_length, [B]
+                paddle.static.InputSpec(
+                    shape=[None, None, None], dtype='float32'),
+                paddle.static.InputSpec(
+                    shape=[None, None, None], dtype='float32')
+            ])
+        return static_model

examples/aishell/s0/local/export.sh

 #!/bin/bash
 
-if [ $# != 3 ];then
-    echo "usage: $0 config_path ckpt_prefix jit_model_path"
+if [ $# != 4 ];then
+    echo "usage: $0 config_path ckpt_prefix jit_model_path model_type"
     exit -1
 fi
 
@@ -11,6 +11,7 @@ echo "using $ngpu gpus..."
 config_path=$1
 ckpt_path_prefix=$2
 jit_model_export_path=$3
+model_type=$4
 
 device=gpu
 if [ ${ngpu} == 0 ];then
@@ -22,8 +23,8 @@ python3 -u ${BIN_DIR}/export.py \
 --nproc ${ngpu} \
 --config ${config_path} \
 --checkpoint_path ${ckpt_path_prefix} \
---export_path ${jit_model_export_path}
-
+--export_path ${jit_model_export_path} \
+--model_type ${model_type}
 
 if [ $? -ne 0 ]; then
     echo "Failed in export!"

examples/aishell/s0/local/test.sh

 #!/bin/bash
 
-if [ $# != 2 ];then
-    echo "usage: ${0} config_path ckpt_path_prefix"
+if [ $# != 3 ];then
+    echo "usage: ${0} config_path ckpt_path_prefix model_type"
     exit -1
 fi
 
@@ -14,6 +14,7 @@ if [ ${ngpu} == 0 ];then
 fi
 config_path=$1
 ckpt_prefix=$2
+model_type=$3
 
 # download language model
 bash local/download_lm_ch.sh
@@ -26,7 +27,8 @@ python3 -u ${BIN_DIR}/test.py \
 --nproc 1 \
 --config ${config_path} \
 --result_file ${ckpt_prefix}.rsl \
---checkpoint_path ${ckpt_prefix}
+--checkpoint_path ${ckpt_prefix} \
+--model_type ${model_type}
 
 if [ $? -ne 0 ]; then
     echo "Failed in evaluation!"

examples/aishell/s0/local/train.sh

 #!/bin/bash
 
-if [ $# != 2 ];then
-    echo "usage: CUDA_VISIBLE_DEVICES=0 ${0} config_path ckpt_name"
+if [ $# != 3 ];then
+    echo "usage: CUDA_VISIBLE_DEVICES=0 ${0} config_path ckpt_name model_type"
     exit -1
 fi
 
@@ -10,6 +10,7 @@ echo "using $ngpu gpus..."
 
 config_path=$1
 ckpt_name=$2
+model_type=$3
 
 device=gpu
 if [ ${ngpu} == 0 ];then
@@ -22,7 +23,8 @@ python3 -u ${BIN_DIR}/train.py \
 --device ${device} \
 --nproc ${ngpu} \
 --config ${config_path} \
---output exp/${ckpt_name}
+--output exp/${ckpt_name} \
+--model_type ${model_type}
 
 if [ $? -ne 0 ]; then
     echo "Failed in training!"

examples/aishell/s0/run.sh

@@ -7,6 +7,7 @@ stage=0
 stop_stage=100
 conf_path=conf/deepspeech2.yaml
 avg_num=1
+model_type=offline
 
 source ${MAIN_ROOT}/utils/parse_options.sh || exit 1;
 
@@ -21,7 +22,7 @@ fi
 
 if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
     # train model, all `ckpt` under `exp` dir
-    CUDA_VISIBLE_DEVICES=${gpus} ./local/train.sh ${conf_path}  ${ckpt}
+    CUDA_VISIBLE_DEVICES=${gpus} ./local/train.sh ${conf_path}  ${ckpt} ${model_type}
 fi
 
 if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
@@ -31,10 +32,10 @@ fi
 
 if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
     # test ckpt avg_n
-    CUDA_VISIBLE_DEVICES=0 ./local/test.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} || exit -1
+    CUDA_VISIBLE_DEVICES=0 ./local/test.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} ${model_type}|| exit -1
 fi
 
 if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
     # export ckpt avg_n
-    CUDA_VISIBLE_DEVICES=0 ./local/export.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} exp/${ckpt}/checkpoints/${avg_ckpt}.jit
+    CUDA_VISIBLE_DEVICES=0 ./local/export.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} exp/${ckpt}/checkpoints/${avg_ckpt}.jit ${model_type}
 fi

examples/librispeech/s0/local/export.sh

 #!/bin/bash
 
-if [ $# != 3 ];then
-    echo "usage: $0 config_path ckpt_prefix jit_model_path"
+if [ $# != 4 ];then
+    echo "usage: $0 config_path ckpt_prefix jit_model_path model_type"
     exit -1
 fi
 
@@ -11,6 +11,7 @@ echo "using $ngpu gpus..."
 config_path=$1
 ckpt_path_prefix=$2
 jit_model_export_path=$3
+model_type=$4
 
 device=gpu
 if [ ${ngpu} == 0 ];then
@@ -22,8 +23,8 @@ python3 -u ${BIN_DIR}/export.py \
 --nproc ${ngpu} \
 --config ${config_path} \
 --checkpoint_path ${ckpt_path_prefix} \
---export_path ${jit_model_export_path}
-
+--export_path ${jit_model_export_path} \
+--model_type ${model_type}
 
 if [ $? -ne 0 ]; then
     echo "Failed in export!"

examples/librispeech/s0/local/test.sh

 #!/bin/bash
 
-if [ $# != 2 ];then
-    echo "usage: ${0} config_path ckpt_path_prefix"
+if [ $# != 3 ];then
+    echo "usage: ${0} config_path ckpt_path_prefix model_type"
     exit -1
 fi
 
@@ -14,6 +14,7 @@ if [ ${ngpu} == 0 ];then
 fi
 config_path=$1
 ckpt_prefix=$2
+model_type=$3
 
 # download language model
 bash local/download_lm_en.sh
@@ -26,7 +27,8 @@ python3 -u ${BIN_DIR}/test.py \
 --nproc 1 \
 --config ${config_path} \
 --result_file ${ckpt_prefix}.rsl \
---checkpoint_path ${ckpt_prefix}
+--checkpoint_path ${ckpt_prefix} \
+--model_type ${model_type}
 
 if [ $? -ne 0 ]; then
     echo "Failed in evaluation!"

examples/librispeech/s0/local/train.sh

 #!/bin/bash
 
-if [ $# != 2 ];then
-    echo "usage: CUDA_VISIBLE_DEVICES=0 ${0} config_path ckpt_name"
+if [ $# != 3 ];then
+    echo "usage: CUDA_VISIBLE_DEVICES=0 ${0} config_path ckpt_name model_type"
     exit -1
 fi
 
@@ -10,6 +10,7 @@ echo "using $ngpu gpus..."
 
 config_path=$1
 ckpt_name=$2
+model_type=$3
 
 device=gpu
 if [ ${ngpu} == 0 ];then
@@ -23,7 +24,8 @@ python3 -u ${BIN_DIR}/train.py \
 --device ${device} \
 --nproc ${ngpu} \
 --config ${config_path} \
---output exp/${ckpt_name}
+--output exp/${ckpt_name} \
+--model_type ${model_type}
 
 if [ $? -ne 0 ]; then
     echo "Failed in training!"

examples/librispeech/s0/run.sh

@@ -6,6 +6,7 @@ stage=0
 stop_stage=100
 conf_path=conf/deepspeech2.yaml
 avg_num=30
+model_type=offline
 source ${MAIN_ROOT}/utils/parse_options.sh || exit 1;
 
 avg_ckpt=avg_${avg_num}
@@ -19,7 +20,7 @@ fi
 
 if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
     # train model, all `ckpt` under `exp` dir
-    CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./local/train.sh ${conf_path}  ${ckpt}
+    CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./local/train.sh ${conf_path}  ${ckpt} ${model_type}
 fi
 
 if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
@@ -29,10 +30,10 @@ fi
 
 if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
     # test ckpt avg_n
-    CUDA_VISIBLE_DEVICES=7 ./local/test.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} || exit -1
+    CUDA_VISIBLE_DEVICES=7 ./local/test.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} ${model_type} || exit -1
 fi
 
 if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
     # export ckpt avg_n
-    CUDA_VISIBLE_DEVICES= ./local/export.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} exp/${ckpt}/checkpoints/${avg_ckpt}.jit
+    CUDA_VISIBLE_DEVICES= ./local/export.sh ${conf_path} exp/${ckpt}/checkpoints/${avg_ckpt} exp/${ckpt}/checkpoints/${avg_ckpt}.jit ${model_type}
 fi