export model test

funasr/export/test_onnx.py

@@ -3,13 +3,13 @@ import numpy as np
 
 
 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)

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])
+