Ver código fonte

Optimize memory (#612)

Co-authored-by: 凌匀 <ailsa.zly@alibaba-inc.com>
lingyunfly 2 anos atrás
pai
commit
6bdae5723e
1 arquivos alterados com 14 adições e 5 exclusões
  1. 14 5
      funasr/models/e2e_vad.py

+ 14 - 5
funasr/models/e2e_vad.py

@@ -252,8 +252,8 @@ class E2EVadModel(nn.Module):
         self.data_buf = None
         self.data_buf_all = None
         self.waveform = None
-        self.ResetDetection()
         self.frontend = frontend
+        self.last_drop_frames = 0
 
     def AllResetDetection(self):
         self.data_buf_start_frame = 0
@@ -282,7 +282,8 @@ class E2EVadModel(nn.Module):
         self.data_buf = None
         self.data_buf_all = None
         self.waveform = None
-        self.ResetDetection()
+        self.last_drop_frames = 0
+        self.windows_detector.Reset()
 
     def ResetDetection(self):
         self.continous_silence_frame_count = 0
@@ -295,6 +296,14 @@ class E2EVadModel(nn.Module):
         self.sil_frame = 0
         self.frame_probs = []
 
+        assert self.output_data_buf[-1].contain_seg_end_point == True
+        drop_frames = int(self.output_data_buf[-1].end_ms / self.vad_opts.frame_in_ms)
+        real_drop_frames = drop_frames - self.last_drop_frames
+        self.last_drop_frames = drop_frames
+        self.data_buf_all = self.data_buf_all[real_drop_frames * int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000):]
+        self.decibel = self.decibel[real_drop_frames:]
+        self.scores = self.scores[:, real_drop_frames:, :]
+
     def ComputeDecibel(self) -> None:
         frame_sample_length = int(self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000)
         frame_shift_length = int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000)
@@ -322,7 +331,7 @@ class E2EVadModel(nn.Module):
         while self.data_buf_start_frame < frame_idx:
             if len(self.data_buf) >= int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000):
                 self.data_buf_start_frame += 1
-                self.data_buf = self.data_buf_all[self.data_buf_start_frame * int(
+                self.data_buf = self.data_buf_all[(self.data_buf_start_frame - self.last_drop_frames) * int(
                     self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000):]
 
     def PopDataToOutputBuf(self, start_frm: int, frm_cnt: int, first_frm_is_start_point: bool,
@@ -543,7 +552,7 @@ class E2EVadModel(nn.Module):
             return 0
         for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1):
             frame_state = FrameState.kFrameStateInvalid
-            frame_state = self.GetFrameState(self.frm_cnt - 1 - i)
+            frame_state = self.GetFrameState(self.frm_cnt - 1 - i - self.last_drop_frames)
             self.DetectOneFrame(frame_state, self.frm_cnt - 1 - i, False)
 
         return 0
@@ -553,7 +562,7 @@ class E2EVadModel(nn.Module):
             return 0
         for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1):
             frame_state = FrameState.kFrameStateInvalid
-            frame_state = self.GetFrameState(self.frm_cnt - 1 - i)
+            frame_state = self.GetFrameState(self.frm_cnt - 1 - i - self.last_drop_frames)
             if i != 0:
                 self.DetectOneFrame(frame_state, self.frm_cnt - 1 - i, False)
             else: