From d567f1922c6d40da0175505cbc9d433f0c2f4cb9 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期二, 07 二月 2023 22:57:08 +0800
Subject: [PATCH] export model test
---
funasr/export/test_onnx.py | 4 ++--
funasr/export/test_torchscripts.py | 17 +++++++++++++++++
2 files changed, 19 insertions(+), 2 deletions(-)
diff --git a/funasr/export/test_onnx.py b/funasr/export/test_onnx.py
index c62137e..4351728 100644
--- a/funasr/export/test_onnx.py
+++ b/funasr/export/test_onnx.py
@@ -3,13 +3,13 @@
if __name__ == '__main__':
- onnx_path = "/Users/zhifu/Downloads/model.onnx"
+ onnx_path = "/mnt/workspace/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/model.onnx"
sess = onnxruntime.InferenceSession(onnx_path)
input_name = [nd.name for nd in sess.get_inputs()]
output_name = [nd.name for nd in sess.get_outputs()]
def _get_feed_dict(feats_length):
- return {'speech': np.zeros((1, feats_length, 560), dtype=np.float32), 'speech_lengths': np.array([feats_length,], dtype=np.int64)}
+ return {'speech': np.zeros((1, feats_length, 560), dtype=np.float32), 'speech_lengths': np.array([feats_length,], dtype=np.int32)}
def _run(feed_dict):
output = sess.run(output_name, input_feed=feed_dict)
diff --git a/funasr/export/test_torchscripts.py b/funasr/export/test_torchscripts.py
new file mode 100644
index 0000000..11be763
--- /dev/null
+++ b/funasr/export/test_torchscripts.py
@@ -0,0 +1,17 @@
+import torch
+import numpy as np
+
+if __name__ == '__main__':
+ onnx_path = "/mnt/workspace/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/model.torchscripts"
+ loaded = torch.jit.load(onnx_path)
+
+ x = torch.rand([2, 21, 560])
+ x_len = torch.IntTensor([6, 21])
+ res = loaded(x, x_len)
+ print(res[0].size(), res[1])
+
+ x = torch.rand([5, 50, 560])
+ x_len = torch.IntTensor([6, 21, 10, 30, 50])
+ res = loaded(x, x_len)
+ print(res[0].size(), res[1])
+
\ No newline at end of file
--
Gitblit v1.9.1