From 38de2af5bf9976d2f14f087d9a0d31991daf6783 Mon Sep 17 00:00:00 2001
From: Zhihao Du <neo.dzh@alibaba-inc.com>
Date: 星期四, 16 三月 2023 19:41:34 +0800
Subject: [PATCH] Merge branch 'main' into dev_dzh

---
 egs_modelscope/asr/paraformer/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/infer.py |   31 ++++++++++++++++++++++---------
 1 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/egs_modelscope/asr/paraformer/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/infer.py b/egs_modelscope/asr/paraformer/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/infer.py
index f9f6114..795a1e7 100644
--- a/egs_modelscope/asr/paraformer/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/infer.py
+++ b/egs_modelscope/asr/paraformer/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/infer.py
@@ -8,9 +8,14 @@
 from funasr.utils.compute_wer import compute_wer
 
 
-def modelscope_infer_core(output_dir, split_dir, njob, idx):
+def modelscope_infer_core(output_dir, split_dir, njob, idx, batch_size, ngpu, model):
     output_dir_job = os.path.join(output_dir, "output.{}".format(idx))
-    gpu_id = (int(idx) - 1) // njob
+    if ngpu > 0:
+        use_gpu = 1
+        gpu_id = int(idx) - 1
+    else:
+        use_gpu = 0
+        gpu_id = -1
     if "CUDA_VISIBLE_DEVICES" in os.environ.keys():
         gpu_list = os.environ['CUDA_VISIBLE_DEVICES'].split(",")
         os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_list[gpu_id])
@@ -18,9 +23,10 @@
         os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
     inference_pipline = pipeline(
         task=Tasks.auto_speech_recognition,
-        model="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
+        model=model,
         output_dir=output_dir_job,
-        batch_size=64
+        batch_size=batch_size,
+        ngpu=use_gpu,
     )
     audio_in = os.path.join(split_dir, "wav.{}.scp".format(idx))
     inference_pipline(audio_in=audio_in)
@@ -30,13 +36,18 @@
     # prepare for multi-GPU decoding
     ngpu = params["ngpu"]
     njob = params["njob"]
+    batch_size = params["batch_size"]
     output_dir = params["output_dir"]
+    model = params["model"]
     if os.path.exists(output_dir):
         shutil.rmtree(output_dir)
     os.mkdir(output_dir)
     split_dir = os.path.join(output_dir, "split")
     os.mkdir(split_dir)
-    nj = ngpu * njob
+    if ngpu > 0:
+        nj = ngpu
+    elif ngpu == 0:
+        nj = njob
     wav_scp_file = os.path.join(params["data_dir"], "wav.scp")
     with open(wav_scp_file) as f:
         lines = f.readlines()
@@ -56,7 +67,7 @@
     p = Pool(nj)
     for i in range(nj):
         p.apply_async(modelscope_infer_core,
-                      args=(output_dir, split_dir, njob, str(i + 1)))
+                      args=(output_dir, split_dir, njob, str(i + 1), batch_size, ngpu, model))
     p.close()
     p.join()
 
@@ -81,8 +92,10 @@
 
 if __name__ == "__main__":
     params = {}
+    params["model"] = "damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"
     params["data_dir"] = "./data/test"
     params["output_dir"] = "./results"
-    params["ngpu"] = 1
-    params["njob"] = 1
-    modelscope_infer(params)
+    params["ngpu"] = 1 # if ngpu > 0, will use gpu decoding
+    params["njob"] = 1 # if ngpu = 0, will use cpu decoding
+    params["batch_size"] = 64
+    modelscope_infer(params)
\ No newline at end of file

--
Gitblit v1.9.1