From 9a9b474e7de7cc90d2ee124dc8d6c2cfa887c059 Mon Sep 17 00:00:00 2001
From: xiaowan0322 <wanchen.swc@alibaba-inc.com>
Date: 星期四, 06 六月 2024 15:59:56 +0800
Subject: [PATCH] [Optimization] support bladedisc fp16 optimization (#1790)

---
 runtime/python/libtorch/funasr_torch/paraformer_bin.py |   26 ++++++++++++++------------
 1 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/runtime/python/libtorch/funasr_torch/paraformer_bin.py b/runtime/python/libtorch/funasr_torch/paraformer_bin.py
index ca96b47..5fa3cc9 100644
--- a/runtime/python/libtorch/funasr_torch/paraformer_bin.py
+++ b/runtime/python/libtorch/funasr_torch/paraformer_bin.py
@@ -282,7 +282,7 @@
                 raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" "\npip3 install -U funasr\n" "For the users in China, you could install with the command:\n" "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple"
 
             model = AutoModel(model=model_dir)
-            model_dir = model.export(type="onnx", quantize=quantize, **kwargs)
+            model_dir = model.export(type="torchscripts", quantize=quantize, **kwargs)
 
         config_file = os.path.join(model_dir, "config.yaml")
         cmvn_file = os.path.join(model_dir, "am.mvn")
@@ -316,9 +316,12 @@
     ) -> List:
         # make hotword list
         hotwords, hotwords_length = self.proc_hotword(hotwords)
-        [bias_embed] = self.eb_infer(torch.Tensor(hotwords), torch.Tensor(hotwords_length))
+        if int(self.device_id) != -1:
+            bias_embed = self.eb_infer(hotwords.cuda())
+        else:
+            bias_embed = self.eb_infer(hotwords)
         # index from bias_embed
-        bias_embed = bias_embed.transpose(1, 0, 2)
+        bias_embed = torch.transpose(bias_embed, 0, 1)
         _ind = np.arange(0, len(hotwords)).tolist()
         bias_embed = bias_embed[_ind, hotwords_length.tolist()]
         waveform_list = self.load_data(wav_content, self.frontend.opts.frame_opts.samp_freq)
@@ -327,15 +330,14 @@
         for beg_idx in range(0, waveform_nums, self.batch_size):
             end_idx = min(waveform_nums, beg_idx + self.batch_size)
             feats, feats_len = self.extract_feat(waveform_list[beg_idx:end_idx])
-            bias_embed = np.expand_dims(bias_embed, axis=0)
-            bias_embed = np.repeat(bias_embed, feats.shape[0], axis=0)
+            bias_embed = torch.unsqueeze(bias_embed, 0).repeat(feats.shape[0], 1, 1)
             try:
                 with torch.no_grad():
                     if int(self.device_id) == -1:
-                        outputs = self.bb_infer(feats, feats_len)
+                        outputs = self.bb_infer(feats, feats_len, bias_embed)
                         am_scores, valid_token_lens = outputs[0], outputs[1]
                     else:
-                        outputs = self.bb_infer_infer(feats.cuda(), feats_len.cuda())
+                        outputs = self.bb_infer(feats.cuda(), feats_len.cuda(), bias_embed.cuda())
                         am_scores, valid_token_lens = outputs[0].cpu(), outputs[1].cpu()
             except:
                 # logging.warning(traceback.format_exc())
@@ -370,16 +372,16 @@
         hotword_int = [word_map(i) for i in hotwords]
         hotword_int.append(np.array([1]))
         hotwords = pad_list(hotword_int, pad_value=0, max_len=10)
-        return hotwords, hotwords_length
+        return torch.tensor(hotwords), hotwords_length
 
     def bb_infer(
         self, feats, feats_len, bias_embed
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        outputs = self.ort_infer_bb([feats, feats_len, bias_embed])
+    ):
+        outputs = self.ort_infer_bb(feats, feats_len, bias_embed)
         return outputs
 
-    def eb_infer(self, hotwords, hotwords_length):
-        outputs = self.ort_infer_eb([hotwords, hotwords_length])
+    def eb_infer(self, hotwords):
+        outputs = self.ort_infer_eb(hotwords.long())
         return outputs
 
     def decode(self, am_scores: np.ndarray, token_nums: int) -> List[str]:

--
Gitblit v1.9.1