Optimize RTF and cer

README.md

@@ -6,7 +6,7 @@ license: apache-2.0
 
 小红书ASR AED-L版本在AX650N上的部署，原项目地址为：[https://github.com/FireRedTeam/FireRedASR](https://github.com/FireRedTeam/FireRedASR)
 
-转换后的模型放置在axmodel目录，目前支持中文、英文，最长输入10秒的音频。
+转换后的模型放置在axmodel目录，目前支持中文、英文，最长输入10秒的音频，超过10秒的音频会用VAD切割后推理。
 
 ## 模型转换
 
@@ -50,121 +50,11 @@ pip install axengine-0.1.3-py3-none-any.whl
 conda activate fireredasr
 python test_ax_model.py
 ```
+```hypo_axmodel.txt```包含识别结果
 
-输出结果如下:
-```
-[INFO] Available providers:  ['AxEngineExecutionProvider']
-Namespace(encoder='axmodel/encoder.axmodel', decoder='axmodel/decoder_main.axmodel', cmvn='axmodel/cmvn.ark', dict='axmodel/dict.txt', spm_model='axmodel/train_bpe1000.model', wavlist='wavlist.txt', hypo='hypo_axmodel.txt', beam_size=3, nbest=1, max_len=128)
-[WARNING] Selected provider(s): ['AXCLRTExecutionProvider'] is(are) not available.
-[INFO] Using provider: AxEngineExecutionProvider
-[INFO] Chip type: ChipType.MC50
-[INFO] VNPU type: VNPUType.DISABLED
-[INFO] Engine version: 2.12.0s
-[INFO] Model type: 2 (triple core)
-[INFO] Compiler version: 4.2 9555977e
-load encoder cost 2.764460325241089 seconds
-[WARNING] Selected provider(s): ['AXCLRTExecutionProvider'] is(are) not available.
-[INFO] Using provider: AxEngineExecutionProvider
-[INFO] Model type: 2 (triple core)
-[INFO] Compiler version: 4.2 9555977e
-load decoder_main cost 16.36833119392395 seconds
-[WARNING] Selected provider(s): ['AXCLRTExecutionProvider'] is(are) not available.
-[INFO] Using provider: AxEngineExecutionProvider
-[INFO] Model type: 2 (triple core)
-[INFO] Compiler version: 4.2 9555977e
-load decoder_loop cost 16.194183826446533 seconds
-run encoder take 196.9749927520752ms
-run decoder_main take 130.2931308746338ms
-run decoder_loop take 165.5733585357666ms
-run decoder_loop take 109.67779159545898ms
-run decoder_loop take 101.15742683410645ms
-run decoder_loop take 110.09836196899414ms
-run decoder_loop take 100.29029846191406ms
-run decoder_loop take 109.33351516723633ms
-run decoder_loop take 100.37779808044434ms
-run decoder_loop take 109.72428321838379ms
-run decoder_loop take 100.42023658752441ms
-run decoder_loop take 101.71890258789062ms
-run decoder_loop take 100.09407997131348ms
-run decoder_loop take 110.25619506835938ms
-run decoder_loop take 100.54206848144531ms
-run decoder_loop take 101.93896293640137ms
-['wav/TEST_NET_Y0000000000_-KTKHdZ2fb8_S00000.wav']
-Durations: 1.8
-Transcribe Durations: 2.5527637004852295
-(Real time factor) RTF: 1.4182020558251274
-wav: wav/TEST_NET_Y0000000000_-KTKHdZ2fb8_S00000.wav
-text: 我有的时候说不清楚你们知道吗
-score: -0.9156361222267151
-
-run encoder take 180.2656650543213ms
-run decoder_main take 91.42565727233887ms
-run decoder_loop take 105.18240928649902ms
-run decoder_loop take 100.56614875793457ms
-run decoder_loop take 100.9066104888916ms
-run decoder_loop take 100.9068489074707ms
-run decoder_loop take 102.90265083312988ms
-run decoder_loop take 100.50129890441895ms
-run decoder_loop take 110.12482643127441ms
-run decoder_loop take 100.65031051635742ms
-run decoder_loop take 110.09883880615234ms
-run decoder_loop take 105.48877716064453ms
-run decoder_loop take 100.32439231872559ms
-run decoder_loop take 106.08601570129395ms
-run decoder_loop take 100.79813003540039ms
-run decoder_loop take 100.4643440246582ms
-run decoder_loop take 100.30460357666016ms
-['wav/TEST_MEETING_T0000000001_S00000.wav']
-Durations: 12.369
-Transcribe Durations: 2.464834690093994
-(Real time factor) RTF: 0.19927517908432324
-wav: wav/TEST_MEETING_T0000000001_S00000.wav
-text: 好首先说一下刚才这个
-score: -0.5064160823822021
-
-run encoder take 172.59907722473145ms
-run decoder_main take 91.79949760437012ms
-run decoder_loop take 105.04364967346191ms
-run decoder_loop take 100.62885284423828ms
-run decoder_loop take 101.89318656921387ms
-run decoder_loop take 100.42643547058105ms
-run decoder_loop take 109.7562313079834ms
-['wav/IT0011W0001.wav']
-Durations: 1.992
-Transcribe Durations: 1.0302071571350098
-(Real time factor) RTF: 0.5171722676380571
-wav: wav/IT0011W0001.wav
-text: 换一首歌
-score: -0.016501454636454582
-
-run encoder take 173.07257652282715ms
-run decoder_main take 91.48693084716797ms
-run decoder_loop take 105.42607307434082ms
-run decoder_loop take 100.10981559753418ms
-run decoder_loop take 100.4478931427002ms
-run decoder_loop take 100.23713111877441ms
-run decoder_loop take 100.10337829589844ms
-run decoder_loop take 100.29196739196777ms
-run decoder_loop take 101.7463207244873ms
-run decoder_loop take 100.8148193359375ms
-run decoder_loop take 109.99274253845215ms
-run decoder_loop take 105.45015335083008ms
-run decoder_loop take 100.59380531311035ms
-run decoder_loop take 100.73733329772949ms
-run decoder_loop take 100.4335880279541ms
-run decoder_loop take 109.68661308288574ms
-['wav/BAC009S0764W0121.wav']
-Durations: 4.2039375
-Transcribe Durations: 2.3024709224700928
-(Real time factor) RTF: 0.5476938994621334
-wav: wav/BAC009S0764W0121.wav
-text: 甚至出现交易几乎停滞的情况
-score: -0.11461181938648224
-
-total wav durations: 20.364937500000003
-total transcribe durations: 8.350276470184326
-AVG RTF: 0.4100320204854213
+## 性能表现
 
-```
+RTF ~= 0.3  
+
+CER(on custom dataset): 3.45%
 
-```hypo_axmodel.txt```包含识别结果

axmodel/decoder_loop.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7c9b3e351557d20846f50d819e18c59d6f10a8adfc40322e5e3034b404b3e038
-size 435136795
+oid sha256:a2912496e6804027f2dc77c903f6b2f76678603dd616e662b78e3f226bcaa91a
+size 416269694

fireredasr/data/asr_feat.py

@@ -42,7 +42,7 @@ class ASRFeatExtractor:
 
         lengths = torch.tensor([feat.size(0) for feat in feats]).long()
         feats_pad = self.pad_feat(feats, 0.0)
-        return feats_pad, lengths, dur
+        return feats_pad.numpy(), lengths, dur
 
     def pad_feat(self, xs, pad_value):
         # type: (List[Tensor], int) -> Tensor

fireredasr_axmodel.py

@@ -10,6 +10,7 @@ from typing import Tuple, List, Dict
 import os
 import time
 import torchaudio
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 try:
     torchaudio.set_audio_backend("soundfile")
@@ -44,18 +45,30 @@ def set_finished_beam_y_to_eos(ys, is_finished, eos_id):
     return ys * (1 - is_finished) + eos_id * is_finished
 
 
+def expand_for_beam_search(n_layer_cross_k, beam_size):
+    """方法1: 使用expand_dims + tile + reshape (最快)"""
+    num_layer, batch_size, Ti, encoder_out_dim = n_layer_cross_k.shape
+    
+    # 在第2维插入新维度
+    expanded = np.expand_dims(n_layer_cross_k, axis=2)
+    # 使用tile替代repeat，性能更好
+    tiled = np.tile(expanded, (1, 1, beam_size, 1, 1))
+    # 重塑形状
+    reshaped = tiled.reshape(num_layer, beam_size * batch_size, Ti, encoder_out_dim)
+    
+    return reshaped
+
+
 class FireRedASRAxModel:
-    def __init__(
-        self,
-        encoder_path: str,
-        decoder_loop_path: str,
-        cmvn_file: str,
-        dict_file: str,
-        spm_model_path: str,
-        providers=["AxEngineExecutionProvider"],
-        decode_max_len=128,
-        audio_dur=10,
-    ):
+    def __init__(self,
+            encoder_path: str,
+            decoder_loop_path: str,
+            cmvn_file: str,
+            dict_file: str,
+            spm_model_path: str,
+            providers=["AxEngineExecutionProvider"],
+            decode_max_len=128,
+            audio_dur=10):
         # NOTE: 参考whisper设置的最大的解码长度
         # FireRedASR-AED 模型支持的最长语音为 60s
         # ref: https://github.com/FireRedTeam/FireRedASR?tab=readme-ov-file#input-length-limitations
@@ -79,6 +92,21 @@ class FireRedASRAxModel:
 
         self.vad_model = load_silero_vad()
 
+        # 预分配内存
+        self._preallocated_memory()
+        # 启用CUDA如果可用
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # print(f"Using device: {self.device}")
+
+    def calc_feat_len(self, audio_dur):
+        import math
+
+        sample_rate = self.sample_rate
+        frame_length = 25 * sample_rate / 1000
+        frame_shift = 10 * sample_rate / 1000
+        length = math.floor((audio_dur * sample_rate - frame_length) / frame_shift) + 1
+        return length
+    
     def init_encoder(self, encoder_path, providers=None):
         self.encoder = axe.InferenceSession(encoder_path, providers=providers)
 
@@ -90,7 +118,7 @@ class FireRedASRAxModel:
         decoder_path = os.path.join(decoder_path, "pe.npy")
 
         return np.load(decoder_path)
-
+    
     def run_encoder(
         self, input: np.ndarray, input_length: np.ndarray
     ) -> Tuple[Tensor, Tensor, Tensor]:
@@ -98,7 +126,7 @@ class FireRedASRAxModel:
             None, {"encoder_input": input, "encoder_input_lengths": input_length}
         )
         return (n_layer_cross_k, n_layer_cross_v, cross_attn_mask)
-
+    
     def decode_loop_one_token(
         self,
         tokens: np.ndarray,
@@ -128,271 +156,485 @@ class FireRedASRAxModel:
             },
         )
         return (logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache)
-
-    def run_decoder(
-        self, n_layer_cross_k, n_layer_cross_v, cross_attn_mask, beam_size, nbest
-    ):
-
+        
+    def _preallocated_memory(self):
+        """预分配常用内存空间"""
+        # 预计算self_attn_mask模板
+        self.self_attn_mask_templates = {}
+        for offset in range(self.decode_max_len):
+            mask = np.zeros((1, 1, self.decode_max_len), dtype=np.float32)
+            mask[:, :, :self.decode_max_len - offset - 1] = -np.inf
+            self.self_attn_mask_templates[offset] = mask
+        
+        # 预分配beam search的scores模板
+        self.beam_scores_template = torch.tensor(
+            [0.0] + [-INF] * (self.decode_max_len - 1)
+        ).float()
+        
+    def transcribe(
+        self, 
+        batch_wav_path: List[str], 
+        beam_size: int = 1, 
+        nbest: int = 1,
+        use_parallel: bool = False
+    ) -> List[Dict]:
+        """优化后的转录方法"""
+        
+        # 1. 优化VAD和分块处理
+        chunks = self._optimized_vad_split(batch_wav_path[0])
+        
+        if use_parallel and len(chunks) > 1:
+            return self._parallel_transcribe(chunks, beam_size, nbest)
+        else:
+            return self._sequential_transcribe(chunks, beam_size, nbest)
+    
+    def _optimized_vad_split(self, wav_path: str) -> List[torch.Tensor]:
+        """优化的VAD分块处理"""
+        import torchaudio
+        
+        # 直接读取为numpy数组，避免torchaudio开销
+        try:
+            wav, sr = torchaudio.load(wav_path)
+            if sr != self.sample_rate:
+                wav = torchaudio.functional.resample(wav, sr, self.sample_rate)
+        except:
+            # 使用silero_vad的read_audio作为备选
+            from silero_vad import read_audio
+            wav = read_audio(wav_path, sampling_rate=self.sample_rate)
+            wav = wav.unsqueeze(0)
+        
+        wav = wav.squeeze(0)
+        
+        # 快速VAD：如果音频较短，直接返回
+        max_chunk_samples = int(self.sample_rate * self.audio_dur)
+        if wav.shape[0] < max_chunk_samples:
+            return [wav]
+        
+        # 使用优化的VAD参数
+        speech_timestamps = get_speech_timestamps(
+            wav,
+            self.vad_model,
+            threshold=0.5,  # 提高阈值，减少静音检测
+            min_speech_duration_ms=250,  # 最小语音段
+            min_silence_duration_ms=100,  # 最小静音段
+            return_seconds=False,
+        )
+        
+        # 优化的分块合并算法
+        return self._optimized_collect_chunks(wav, speech_timestamps)
+    
+    def _optimized_collect_chunks(
+        self, 
+        wav: torch.Tensor, 
+        speech_timestamps: List[Dict]
+    ) -> List[torch.Tensor]:
+        """优化的分块合并算法"""
+        max_chunk_samples = int(self.sample_rate * self.audio_dur)
+        chunks = []
+        current_chunk = []
+        current_length = 0
+        
+        for ts in speech_timestamps:
+            start, end = ts["start"], ts["end"]
+            chunk_length = end - start
+            
+            if current_length + chunk_length <= max_chunk_samples:
+                current_chunk.append((start, end))
+                current_length += chunk_length
+            else:
+                if current_chunk:
+                    # 合并当前chunk
+                    merged = torch.cat([wav[s:e] for s, e in current_chunk])
+                    chunks.append(merged)
+                
+                if chunk_length > max_chunk_samples:
+                    # 大chunk分割
+                    num_splits = (chunk_length + max_chunk_samples - 1) // max_chunk_samples
+                    for i in range(num_splits):
+                        s = start + i * max_chunk_samples
+                        e = min(start + (i + 1) * max_chunk_samples, end)
+                        chunks.append(wav[s:e])
+                    current_chunk = []
+                    current_length = 0
+                else:
+                    current_chunk = [(start, end)]
+                    current_length = chunk_length
+        
+        # 处理最后一个chunk
+        if current_chunk:
+            merged = torch.cat([wav[s:e] for s, e in current_chunk])
+            chunks.append(merged)
+        
+        return chunks
+    
+    def _optimized_decode_loop(
+        self, 
+        n_layer_cross_k: np.ndarray,
+        n_layer_cross_v: np.ndarray,
+        cross_attn_mask: np.ndarray,
+        beam_size: int,
+        nbest: int
+    ) -> List[Dict]:
+        """优化的解码循环"""
+        
         num_layer, batch_size, Ti, encoder_out_dim = n_layer_cross_k.shape
         encoder_out_length = cross_attn_mask.shape[-1]
 
-        cross_attn_mask = torch.from_numpy(cross_attn_mask).to(torch.float32)
-        cross_attn_mask = (
-            cross_attn_mask.unsqueeze(1)
-            .repeat(1, beam_size, 1, 1)
-            .view(beam_size * batch_size, -1, encoder_out_length)
-        )
-
-        n_layer_cross_k = torch.from_numpy(n_layer_cross_k)
-        n_layer_cross_v = torch.from_numpy(n_layer_cross_v)
-        n_layer_cross_k = (
-            n_layer_cross_k.unsqueeze(2)
-            .repeat(1, 1, beam_size, 1, 1)
-            .view(num_layer, beam_size * batch_size, Ti, encoder_out_dim)
-        )
-        n_layer_cross_v = (
-            n_layer_cross_v.unsqueeze(2)
-            .repeat(1, 1, beam_size, 1, 1)
-            .view(num_layer, beam_size * batch_size, Ti, encoder_out_dim)
-        )
-
-        prediction_tokens = (
-            torch.ones(beam_size * batch_size, 1).fill_(self.sos_id).long()
-        )
-        tokens = prediction_tokens
-        offset = torch.zeros(1, dtype=torch.int64)
-        n_layer_self_k_cache, n_layer_self_v_cache = self.get_initialized_self_cache(
+        n_layer_cross_k = expand_for_beam_search(n_layer_cross_k, beam_size)
+        n_layer_cross_v = expand_for_beam_search(n_layer_cross_v, beam_size)
+
+        batch_size, Ti, encoder_out_length = cross_attn_mask.shape
+    
+        # 在第1维插入新维度
+        expanded = np.expand_dims(cross_attn_mask, axis=1)
+        # 使用tile替代repeat，性能更好
+        tiled = np.tile(expanded, (1, beam_size, 1, 1))
+        # 重塑形状
+        cross_attn_mask = tiled.reshape(beam_size * batch_size, Ti, encoder_out_length)
+            
+        # 优化的cache初始化
+        n_layer_self_k_cache, n_layer_self_v_cache = self._optimized_init_self_cache(
             batch_size, beam_size
         )
-
-        scores = torch.tensor([0.0] + [-INF] * (beam_size - 1)).float()
-        scores = scores.repeat(batch_size).view(batch_size * beam_size, 1)
-        is_finished = torch.zeros_like(scores)
-
-        self_attn_mask = np.zeros((batch_size * beam_size, 1, 1), dtype=np.float32)
-
-        for i in range(self.decode_max_len):
-
-            tokens = to_numpy(tokens).astype(np.int32)
-            n_layer_self_k_cache = to_numpy(n_layer_self_k_cache)
-            n_layer_self_v_cache = to_numpy(n_layer_self_v_cache)
-            n_layer_cross_k = to_numpy(n_layer_cross_k)
-            n_layer_cross_v = to_numpy(n_layer_cross_v)
-            cross_attn_mask = to_numpy(cross_attn_mask)
-
-            self_attn_mask = np.zeros(
-                (batch_size * beam_size, 1, self.decode_max_len), dtype=np.float32
+        
+        # 预分配tokens和scores
+        tokens = torch.full(
+            (beam_size * batch_size, 1), 
+            self.sos_id, 
+            dtype=torch.int32, device=self.device
+        )
+        scores = self.beam_scores_template[:beam_size].repeat(batch_size).view(
+            batch_size * beam_size, 1
+        ).to(self.device)
+        is_finished = torch.zeros_like(scores, dtype=torch.bool, device=self.device)
+        
+        # 预分配prediction_tokens
+        prediction_tokens = tokens.clone()
+        
+        pe_np = self.pe
+        
+        for offset in range(self.decode_max_len):
+            # 使用预计算的mask模板
+            self_attn_mask = np.repeat(
+                self.self_attn_mask_templates[offset], 
+                beam_size * batch_size, 
+                axis=0
             )
-            self_attn_mask[:, :, : self.decode_max_len - offset[0] - 1] = -np.inf
-
-            (
-                logits,
-                n_layer_self_k_cache,
-                n_layer_self_v_cache,
-            ) = self.decode_loop_one_token(
-                to_numpy(tokens),
-                to_numpy(n_layer_self_k_cache),
-                to_numpy(n_layer_self_v_cache),
-                to_numpy(n_layer_cross_k),
-                to_numpy(n_layer_cross_v),
-                self.pe[offset],
-                self_attn_mask,
-                to_numpy(cross_attn_mask),
+            
+            # 直接使用numpy数组，避免转换
+            logits, n_layer_self_k_cache, n_layer_self_v_cache = (
+                self.decode_loop_one_token(
+                    tokens.cpu().numpy().astype(np.int32),
+                    n_layer_self_k_cache,
+                    n_layer_self_v_cache,
+                    n_layer_cross_k,
+                    n_layer_cross_v,
+                    pe_np[offset],
+                    self_attn_mask,
+                    cross_attn_mask
+                )
             )
-
-            offset += 1
-            logits = torch.from_numpy(logits)
-
-            logits = logits.squeeze(1)
+            
+            logits = torch.from_numpy(logits).to(self.device).squeeze(1)
             t_scores = F.log_softmax(logits, dim=-1)
-            t_topB_scores, t_topB_ys = torch.topk(t_scores, k=beam_size, dim=1)
-            t_topB_scores = set_finished_beam_score_to_zero(t_topB_scores, is_finished)
-            t_topB_ys = set_finished_beam_y_to_eos(t_topB_ys, is_finished, self.eos_id)
-
-            scores = scores + t_topB_scores
-
-            scores = scores.view(batch_size, beam_size * beam_size)
-            scores, topB_score_ids = torch.topk(scores, k=beam_size, dim=1)
-            scores = scores.view(-1, 1)
-
-            topB_row_number_in_each_B_rows_of_ys = torch.div(
-                topB_score_ids, beam_size
-            ).view(batch_size * beam_size)
-            stride = beam_size * torch.arange(batch_size).view(batch_size, 1).repeat(
-                1, beam_size
-            ).view(batch_size * beam_size)
-            topB_row_number_in_ys = (
-                topB_row_number_in_each_B_rows_of_ys.long() + stride.long()
+            
+            # 优化的beam search
+            tokens, scores, prediction_tokens, n_layer_self_k_cache, n_layer_self_v_cache, is_finished = (
+                self._optimized_beam_search(
+                    t_scores, tokens, scores, prediction_tokens,
+                    n_layer_self_k_cache, n_layer_self_v_cache,
+                    is_finished, beam_size, batch_size
+                )
             )
-
-            prediction_tokens = prediction_tokens[topB_row_number_in_ys]
-            t_ys = torch.gather(
-                t_topB_ys.view(batch_size, beam_size * beam_size),
-                dim=1,
-                index=topB_score_ids,
-            ).view(beam_size * batch_size, 1)
-
-            tokens = t_ys
-
-            prediction_tokens = torch.cat((prediction_tokens, t_ys), dim=1)
-
-            n_layer_self_k_cache = torch.from_numpy(n_layer_self_k_cache)
-            n_layer_self_v_cache = torch.from_numpy(n_layer_self_v_cache)
-
-            for i, self_k_cache in enumerate(n_layer_self_k_cache):
-                n_layer_self_k_cache[i] = n_layer_self_k_cache[i][topB_row_number_in_ys]
-
-            for i, self_v_cache in enumerate(n_layer_self_v_cache):
-                n_layer_self_v_cache[i] = n_layer_self_v_cache[i][topB_row_number_in_ys]
-
-            is_finished = t_ys.eq(self.eos_id)
-            if is_finished.sum().item() == beam_size * batch_size:
+            
+            if is_finished.all():
                 break
-
+        
+        # return self._extract_results(scores, prediction_tokens, batch_size, beam_size, nbest)
+        return self.extract_results_numpy_vectorized(scores.numpy(), prediction_tokens.numpy(), batch_size, beam_size, nbest)
+    
+    
+    def _optimized_beam_search(
+        self,
+        t_scores: torch.Tensor,
+        tokens: torch.Tensor,
+        scores: torch.Tensor,
+        prediction_tokens: torch.Tensor,
+        n_layer_self_k_cache: torch.Tensor,
+        n_layer_self_v_cache: torch.Tensor,
+        is_finished: torch.Tensor,
+        beam_size: int,
+        batch_size: int
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """优化的beam search步骤"""
+        
+        # 使用torch的in-place操作
+        t_topB_scores, t_topB_ys = torch.topk(t_scores, k=beam_size, dim=1)
+        
+        # 处理已完成的beam
+        if is_finished.any():
+            # 原地操作，避免创建新tensor
+            t_topB_scores.masked_fill_(is_finished, 0.0)
+            t_topB_scores[:, 1:].masked_fill_(is_finished, -INF)
+            t_topB_ys.masked_fill_(is_finished, self.eos_id)
+        
+        # 更新scores
+        scores = scores + t_topB_scores
+        
+        # 优化的topk选择
+        scores_2d = scores.view(batch_size, beam_size * beam_size)
+        top_scores, top_ids = torch.topk(scores_2d, k=beam_size, dim=1)
+        scores = top_scores.view(-1, 1)
+        
+        # 计算索引
+        topB_row_number_in_each_B_rows_of_ys = torch.div(top_ids, beam_size, rounding_mode='floor')
+        stride = beam_size * torch.arange(batch_size, device=self.device).view(batch_size, 1)
+        topB_row_number_in_ys = (topB_row_number_in_each_B_rows_of_ys + stride).view(-1)
+        
+        # 更新tokens和prediction_tokens
+        tokens = torch.gather(
+            t_topB_ys.view(batch_size, beam_size * beam_size),
+            dim=1,
+            index=top_ids,
+        ).view(beam_size * batch_size, 1)
+        
+        prediction_tokens = torch.cat([
+            prediction_tokens[topB_row_number_in_ys],
+            tokens
+        ], dim=1)
+        
+        # 更新cache（原地操作）
+        for i in range(n_layer_self_k_cache.shape[0]):
+            n_layer_self_k_cache[i] = n_layer_self_k_cache[i][topB_row_number_in_ys]
+            n_layer_self_v_cache[i] = n_layer_self_v_cache[i][topB_row_number_in_ys]
+        
+        # 更新完成状态
+        is_finished = tokens.eq(self.eos_id)
+        
+        return tokens, scores, prediction_tokens, n_layer_self_k_cache, n_layer_self_v_cache, is_finished
+    
+    def _optimized_init_self_cache(
+        self, batch_size: int, beam_size: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """优化的self cache初始化"""
+        shape = (
+            self.num_decoder_blocks,
+            batch_size * beam_size,
+            self.decode_max_len,
+            self.decoder_hidden_dim
+        )
+        n_layer_self_k_cache = np.zeros(shape, dtype=np.float32)
+        n_layer_self_v_cache = np.zeros(shape, dtype=np.float32)
+        return n_layer_self_k_cache, n_layer_self_v_cache
+    
+    def _extract_results(
+        self,
+        scores: torch.Tensor,
+        prediction_tokens: torch.Tensor,
+        batch_size: int,
+        beam_size: int,
+        nbest: int
+    ) -> List[Dict]:
+        """提取结果"""
         scores = scores.view(batch_size, beam_size)
-        prediction_valid_token_lengths = torch.sum(
+        valid_lengths = torch.sum(
             torch.ne(prediction_tokens.view(batch_size, beam_size, -1), self.eos_id),
-            dim=-1,
+            dim=-1
         ).int()
-
+        
         nbest_scores, nbest_ids = torch.topk(scores, k=nbest, dim=1)
-        index = (
-            nbest_ids + beam_size * torch.arange(batch_size).view(batch_size, 1).long()
-        )
-        nbest_prediction_tokens = prediction_tokens.view(batch_size * beam_size, -1)[
-            index.view(-1)
-        ]
-        nbest_prediction_tokens = nbest_prediction_tokens.view(
-            batch_size, nbest_ids.size(1), -1
-        )
-        nbest_prediction_valid_token_lengths = prediction_valid_token_lengths.view(
-            batch_size * beam_size
-        )[index.view(-1)].view(batch_size, -1)
-
-        # batch_size is always 1
-        i_best_hyps: List[Dict[str, torch.Tensor]] = []
+        index = nbest_ids + beam_size * torch.arange(batch_size, device=self.device).unsqueeze(1)
+        
+        nbest_tokens = prediction_tokens.view(batch_size * beam_size, -1)[index.view(-1)]
+        nbest_tokens = nbest_tokens.view(batch_size, nbest_ids.size(1), -1)
+        
+        results = []
         for j, score in enumerate(nbest_scores[0]):
             hyp = {
-                "token_ids": nbest_prediction_tokens[
-                    0, j, 1 : nbest_prediction_valid_token_lengths[0, j]
-                ],
+                "token_ids": nbest_tokens[0, j, 1:valid_lengths[0, nbest_ids[0, j]]],
                 "score": score,
             }
-            i_best_hyps.append(hyp)
-
-        return i_best_hyps
-
-    def get_initialized_self_cache(
-        self, batch_size, beam_size
-    ) -> Tuple[Tensor, Tensor]:
-        n_layer_self_k_cache = torch.zeros(
-            self.num_decoder_blocks,
-            batch_size * beam_size,
-            self.decode_max_len,
-            self.decoder_hidden_dim,
-        )
-        n_layer_self_v_cache = torch.zeros(
-            self.num_decoder_blocks,
-            batch_size * beam_size,
-            self.decode_max_len,
-            self.decoder_hidden_dim,
-        )
-        return n_layer_self_k_cache, n_layer_self_v_cache
-
-    def calc_feat_len(self, audio_dur):
-        import math
-
-        sample_rate = self.sample_rate
-        frame_length = 25 * sample_rate / 1000
-        frame_shift = 10 * sample_rate / 1000
-        length = math.floor((audio_dur * sample_rate - frame_length) / frame_shift) + 1
-        return length
-
-    def collect_chunks(self, wav, speech_timestamps, audio_dur, sample_rate):
-        max_chunk_samples = int(audio_dur * sample_rate)
-        chunks = []
-        for ts in speech_timestamps:
-            start, end = ts["start"], ts["end"]
-            cur_chunk = wav[start:end]
-            if (
-                len(chunks) > 0
-                and chunks[-1].shape[0] + cur_chunk.shape[0] < max_chunk_samples
-            ):
-                chunks[-1] = torch.concat([chunks[-1], cur_chunk], dim=0)
-            else:
-                if cur_chunk.shape[0] > max_chunk_samples:
-                    # greedy split if one chunk is too big
-                    chunks.append(cur_chunk[:max_chunk_samples])
-                    chunks.append(cur_chunk[max_chunk_samples:])
-                else:
-                    chunks.append(cur_chunk)
-        return chunks
-
-    def transcribe(
-        self, batch_wav_path: List[str], beam_size: int = 1, nbest: int = 1
+            results.append(hyp)
+        
+        return results
+    
+    
+    def extract_results_numpy_vectorized(
+        self,
+        scores: np.ndarray,
+        prediction_tokens: np.ndarray,
+        batch_size: int,
+        beam_size: int,
+        nbest: int,
+        eos_id: int = 4
     ) -> List[Dict]:
-
-        # Run vad, greedy split audio to fit audio_dur
-        try:
-            wav = read_audio(batch_wav_path[0], sampling_rate=self.sample_rate)
-        except Exception as e:
-            print("Please run apt install libsnffile1 first")
-            raise e
-
-        max_chunk_samples = int(self.sample_rate * self.audio_dur)
-        if wav.shape[0] < max_chunk_samples:
-            chunks = [wav]
-        else:
-            speech_timestamps = get_speech_timestamps(
-                wav,
-                self.vad_model,
-                return_seconds=False,  # Return speech timestamps in seconds (default is samples)
-            )
-            chunks = self.collect_chunks(
-                wav, speech_timestamps, self.audio_dur, self.sample_rate
-            )
-            # print(f"Split to {len(chunks)} chunks")
-
-        transcribe_durations = 0
-        wav_durations = []
+        """向量化版本的NumPy实现"""
+        
+        # 1. 重塑和计算有效长度
+        scores_2d = scores.reshape(batch_size, beam_size)
+        tokens_3d = prediction_tokens.reshape(batch_size, beam_size, -1)
+        
+        # 计算有效长度（不包括eos_id）
+        valid_lengths = np.sum(tokens_3d != eos_id, axis=-1).astype(np.int32)
+        
+        # 2. 使用argpartition进行部分排序（比argsort更快）
+        # 获取最大的nbest个元素的索引
+        # 使用argpartition: O(n) vs argsort: O(n log n)
+        partitioned_indices = np.argpartition(-scores_2d, nbest-1, axis=1)[:, :nbest]
+        
+        # 对每个batch内的topk进行排序
+        nbest_scores = np.take_along_axis(scores_2d, partitioned_indices, axis=1)
+        sorted_order = np.argsort(-nbest_scores, axis=1)
+        
+        # 应用排序
+        nbest_ids = np.take_along_axis(partitioned_indices, sorted_order, axis=1)
+        nbest_scores = np.take_along_axis(nbest_scores, sorted_order, axis=1)
+        
+        # 3. 计算全局索引
+        batch_indices = np.arange(batch_size)[:, np.newaxis]
+        global_indices = nbest_ids + beam_size * batch_indices
+        flat_global_indices = global_indices.reshape(-1)
+        
+        # 4. 提取tokens
+        flat_tokens = prediction_tokens.reshape(-1, prediction_tokens.shape[-1])
+        nbest_tokens = flat_tokens[flat_global_indices]
+        nbest_tokens = nbest_tokens.reshape(batch_size, nbest, -1)
+        
+        # 5. 提取对应的有效长度
+        nbest_valid_lengths = np.take_along_axis(valid_lengths, nbest_ids, axis=1)
+        
+        # 6. 构建结果
+        results = []
+        
+        for b in range(batch_size):
+            batch_results = []
+            for j in range(nbest):
+                valid_len = nbest_valid_lengths[b, j]
+                
+                # 提取token_ids（跳过<sos>）
+                token_ids = nbest_tokens[b, j, 1:valid_len]
+                
+                hyp = {
+                    "token_ids": token_ids.tolist(),
+                    "score": float(nbest_scores[b, j]),
+                }
+                batch_results.append(hyp)
+            
+            # 如果是批量处理，可以按batch返回
+            # 这里假设batch_size=1，直接返回第一个batch的结果
+            if b == 0:
+                results = batch_results
+        
+        return results
+    
+
+    def _sequential_transcribe(
+        self, 
+        chunks: List[torch.Tensor], 
+        beam_size: int, 
+        nbest: int
+    ) -> Dict:
+        """顺序转录（单线程）"""
         tokens = []
+        wav_durations = []
+        transcribe_duration = 0
+        
         for chunk in chunks:
-            chunk = (chunk.clamp(-1, 1) * 32768).to(torch.int16)
-            feats, lengths, wav_duration = self.feature_extractor.run_chunk(
-                chunk, self.sample_rate
-            )
-
+            # 优化的特征提取
+            feats, lengths, wav_duration = self._optimized_feature_extraction(chunk)
             wav_durations.append(wav_duration)
-
-            if feats.shape[1] < self.max_feat_len:
-                feats = np.concatenate(
-                    [
-                        feats,
-                        np.zeros(
-                            (1, self.max_feat_len - feats.shape[1], 80),
-                            dtype=np.float32,
-                        ),
-                    ],
-                    axis=1,
-                )
-            feats = feats[:, : self.max_feat_len, :]
-            lengths = torch.minimum(lengths, torch.tensor(self.max_feat_len))
-
-            feats = to_numpy(feats)
-            lengths = to_numpy(lengths).astype(np.int32)
-
+            
+            # 运行encoder和decoder
             start_time = time.time()
             n_layer_cross_k, n_layer_cross_v, cross_attn_mask = self.run_encoder(
-                to_numpy(feats), to_numpy(lengths)
+                feats, lengths.numpy().astype(np.int32)
             )
-            # print(f"run encoder take {(time.time() - start_time) * 1000}ms")
-            nbest_hyps = self.run_decoder(
+            
+            nbest_hyps = self._optimized_decode_loop(
                 n_layer_cross_k, n_layer_cross_v, cross_attn_mask, beam_size, nbest
             )
-            tokens.extend([int(id) for id in nbest_hyps[0]["token_ids"].cpu()])
-
-            transcribe_durations += time.time() - start_time
-
+            
+            tokens.extend([int(id) for id in nbest_hyps[0]["token_ids"]])
+            transcribe_duration += time.time() - start_time
+        
         text = self.tokenizer.detokenize(tokens)
-
-        return {"text": text}, wav_durations, transcribe_durations
+        return {"text": text}, wav_durations, transcribe_duration
+    
+    def _parallel_transcribe(
+        self, 
+        chunks: List[torch.Tensor], 
+        beam_size: int, 
+        nbest: int
+    ) -> Dict:
+        """并行转录（多线程）"""
+        import threading
+        
+        results = []
+        lock = threading.Lock()
+        
+        def process_chunk(chunk_idx, chunk):
+            try:
+                # 特征提取
+                feats, lengths, wav_duration = self._optimized_feature_extraction(chunk)
+                
+                # encoder
+                n_layer_cross_k, n_layer_cross_v, cross_attn_mask = self.run_encoder(
+                    feats, lengths.astype(np.int32)
+                )
+                
+                # decoder
+                nbest_hyps = self._optimized_decode_loop(
+                    n_layer_cross_k, n_layer_cross_v, cross_attn_mask, beam_size, nbest
+                )
+                
+                with lock:
+                    results.append({
+                        'chunk_idx': chunk_idx,
+                        'tokens': [int(id) for id in nbest_hyps[0]["token_ids"].cpu()],
+                        'duration': wav_duration
+                    })
+            except Exception as e:
+                print(f"Error processing chunk {chunk_idx}: {e}")
+        
+        # 使用ThreadPoolExecutor并行处理
+        with ThreadPoolExecutor(max_workers=min(4, len(chunks))) as executor:
+            futures = []
+            for i, chunk in enumerate(chunks):
+                future = executor.submit(process_chunk, i, chunk)
+                futures.append(future)
+            
+            # 等待所有任务完成
+            for future in as_completed(futures):
+                future.result()
+        
+        # 合并结果
+        results.sort(key=lambda x: x['chunk_idx'])
+        tokens = []
+        wav_durations = []
+        
+        for result in results:
+            tokens.extend(result['tokens'])
+            wav_durations.append(result['duration'])
+        
+        text = self.tokenizer.detokenize(tokens)
+        return {"text": text}, wav_durations, 0  # 并行处理时间不好统计
+    
+    def _optimized_feature_extraction(
+        self, 
+        chunk: torch.Tensor
+    ) -> Tuple[np.ndarray, np.ndarray, float]:
+        """优化的特征提取"""
+        chunk = (chunk.clamp(-1, 1) * 32768).to(torch.int16)
+        feats, lengths, wav_duration = self.feature_extractor.run_chunk(
+            chunk, self.sample_rate
+        )
+        
+        # 原地padding，避免创建新数组
+        if feats.shape[1] < self.max_feat_len:
+            pad_width = ((0, 0), (0, self.max_feat_len - feats.shape[1]), (0, 0))
+            feats = np.pad(feats, pad_width, mode='constant', constant_values=0)
+        
+        feats = feats[:, :self.max_feat_len, :]
+        lengths = np.minimum(lengths, self.max_feat_len)
+        
+        return feats, lengths, wav_duration

test_ax_model.py

@@ -44,7 +44,7 @@ def parse_args():
     parser.add_argument(
         "--hypo", type=str, default="hypo_axmodel.txt", help="File of hypos"
     )
-    parser.add_argument("--beam_size", type=int, default=3, help="")
+    parser.add_argument("--beam_size", type=int, default=1, help="")
     parser.add_argument("--nbest", type=int, default=1, help="")
     parser.add_argument("--decode_max_len", type=int, default=128, help="max token len")
     parser.add_argument("--max_dur", type=int, default=10, help="max audio len")

test_wer.py

@@ -183,12 +183,6 @@ def get_args():
         default="axmodel/encoder.axmodel",
         help="Path to onnx encoder",
     )
-    parser.add_argument(
-        "--decoder_main",
-        type=str,
-        default="axmodel/decoder_main.axmodel",
-        help="Path to axmodel decoder main",
-    )
     parser.add_argument(
         "--decoder_loop",
         type=str,
@@ -213,7 +207,7 @@ def get_args():
     parser.add_argument(
         "--hypo", type=str, default="hypo_axmodel.txt", help="File of hypos"
     )
-    parser.add_argument("--beam_size", type=int, default=3, help="")
+    parser.add_argument("--beam_size", type=int, default=1, help="")
     parser.add_argument("--nbest", type=int, default=1, help="")
     parser.add_argument("--max_len", type=int, default=128, help="")
     return parser.parse_args()