From 4a7a984a5f3e3f894f86ce82e76ddd13d8a42a20 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期一, 11 三月 2024 17:56:30 +0800
Subject: [PATCH] Dev gzf (#1465)
---
funasr/models/fsmn_vad_streaming/model.py | 45 +++++++++++++++++++++++++++++++++++++++++++++
1 files changed, 45 insertions(+), 0 deletions(-)
diff --git a/funasr/models/fsmn_vad_streaming/model.py b/funasr/models/fsmn_vad_streaming/model.py
index e29af57..c3063b0 100644
--- a/funasr/models/fsmn_vad_streaming/model.py
+++ b/funasr/models/fsmn_vad_streaming/model.py
@@ -284,6 +284,7 @@
encoder_class = tables.encoder_classes.get(encoder)
encoder = encoder_class(**encoder_conf)
self.encoder = encoder
+ self.encoder_conf = encoder_conf
def ResetDetection(self, cache: dict = {}):
cache["stats"].continous_silence_frame_count = 0
@@ -642,6 +643,50 @@
return results, meta_data
+ def export(self, **kwargs):
+ is_onnx = kwargs.get("type", "onnx") == "onnx"
+ encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export")
+ self.encoder = encoder_class(self.encoder, onnx=is_onnx)
+ self.forward = self._export_forward
+
+ return self
+
+ def _export_forward(self, feats: torch.Tensor, *args, **kwargs):
+
+ scores, out_caches = self.encoder(feats, *args)
+
+ return scores, out_caches
+
+ def export_dummy_inputs(self, data_in=None, frame=30):
+ if data_in is None:
+ speech = torch.randn(1, frame, self.encoder_conf.get("input_dim"))
+ else:
+ speech = None # Undo
+
+ cache_frames = self.encoder_conf.get("lorder") + self.encoder_conf.get("rorder") - 1
+ in_cache0 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1)
+ in_cache1 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1)
+ in_cache2 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1)
+ in_cache3 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1)
+
+ return (speech, in_cache0, in_cache1, in_cache2, in_cache3)
+
+ def export_input_names(self):
+ return ['speech', 'in_cache0', 'in_cache1', 'in_cache2', 'in_cache3']
+
+ def export_output_names(self):
+ return ['logits', 'out_cache0', 'out_cache1', 'out_cache2', 'out_cache3']
+
+ def export_dynamic_axes(self):
+ return {
+ 'speech': {
+ 1: 'feats_length'
+ },
+ }
+
+ def export_name(self, ):
+ return "model.onnx"
+
def DetectCommonFrames(self, cache: dict = {}) -> int:
if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateEndPointDetected:
return 0
--
Gitblit v1.9.1