logo
Browse Source

Update onnx test

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 1 year ago
parent
commit
36653e1a14
  1. 5
      test_onnx.py
  2. 4
      test_onnx2.py

5
test_onnx.py

@ -1,5 +1,4 @@
from auto_transformers import AutoTransformers
from towhee import ops
import numpy
import onnx
import onnxruntime
@ -48,7 +47,7 @@ for name in models:
f.write(','.join(status) + '\n')
status = [name] + ['fail'] * 5
try:
op = AutoTransformers(model_name=name)
op = ops.text_embedding.transformers(model_name=name, device='cpu').get_op()
out1 = op(test_txt)
logger.info('OP LOADED.')
status[1] = 'success'

4
test_onnx2.py

@ -1,4 +1,4 @@
from auto_transformers import AutoTransformers
from towhee import ops
import torch
import numpy
import onnx
@ -82,7 +82,7 @@ for name in models:
f.write(','.join(status) + '\n')
status = [name] + ['fail'] * 5
try:
op = AutoTransformers(model_name=name, device='cpu')
op = ops.text_embedding.transformers(model_name=name, device='cpu').get_op()
out1 = op(test_txt)
logger.info('OP LOADED.')
status[1] = 'success'

Loading…
Cancel
Save