export model test

This commit is contained in:
游雁 2023-02-07 22:57:08 +08:00
parent 87bff7ae59
commit d567f1922c
2 changed files with 19 additions and 2 deletions

View File

@ -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)

View File

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