logo
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions

37 lines
1.1 KiB

from auto_transformers import AutoTransformers
import onnx
f = open('onnx.csv', 'a+')
f.write('model_name, run_op, save_onnx, check_onnx\n')
# models = AutoTransformers.supported_model_names()[:1]
models = ['bert-base-cased', 'distilbert-base-cased']
for name in models:
f.write(f'{name},')
try:
op = AutoTransformers(model_name=name)
out1 = op('hello, world.')
f.write('success,')
except Exception as e:
f.write('fail')
print(f'Fail to load op for {name}: {e}')
continue
try:
op.save_model(format='onnx')
f.write('success,')
except Exception as e:
f.write('fail')
print(f'Fail to save onnx for {name}: {e}')
continue
try:
saved_name = name.replace('/', '-')
onnx_model = onnx.load(f'saved/onnx/{saved_name}.onnx')
onnx.checker.check_model(onnx_model)
f.write('success')
except Exception as e:
f.write('fail')
print(f'Fail to check onnx for {name}: {e}')
continue
f.write('\n')
print('Finished.')