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

40 lines
1.1 KiB

from auto_transformers import AutoTransformers
import torch
f = open('torchscript.csv', 'a+')
f.write('model_name, run op, save_torchscript, check_result\n')
models = AutoTransformers.supported_model_names()[:1]
for name in models:
line = f'{name}, '
try:
op = AutoTransformers(model_name=name)
out1 = op('hello, world.')
line += 'success, '
except Exception as e:
line += 'fail\n'
f.write(line)
print(f'Fail to load op for {name}: {e}.')
continue
try:
op.save_model(format='torchscript')
line += 'success, '
except Exception as e:
line += 'fail\n'
f.write(line)
print(f'Fail to save onnx for {name}: {e}.')
continue
try:
saved_name = name.replace('/', '-')
op.model = torch.jit.load(f'saved/torchscript/{saved_name}.pt')
out2 = op('hello, world.')
assert (out1 == out2).all()
line += 'success'
except Exception as e:
line += 'fail\n'
f.write(line)
print(f'Fail to check onnx for {name}: {e}.')
continue
line += '\n'
f.write(line)