From d567f1922c6d40da0175505cbc9d433f0c2f4cb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B8=B8=E9=9B=81?= Date: Tue, 7 Feb 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(-) create mode 100644 funasr/export/test_torchscripts.py diff --git a/funasr/export/test_onnx.py b/funasr/export/test_onnx.py index c62137ea7..435172850 100644 --- a/funasr/export/test_onnx.py +++ b/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) diff --git a/funasr/export/test_torchscripts.py b/funasr/export/test_torchscripts.py new file mode 100644 index 000000000..11be76325 --- /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