From 9acb9072e13dcb3d6cea8facb452d6bd15efcfe0 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Wed, 7 Dec 2022 15:13:16 +0800 Subject: [PATCH] Fix log Signed-off-by: Jael Gu --- auto_transformers.py | 44 +++++++++++++++++++++++++------------------- test_onnx.py | 30 +++++++++++++++++++++++------- 2 files changed, 48 insertions(+), 26 deletions(-) diff --git a/auto_transformers.py b/auto_transformers.py index 7e51693..9ff2488 100644 --- a/auto_transformers.py +++ b/auto_transformers.py @@ -22,6 +22,9 @@ from towhee.operator import NNOperator from towhee import register import warnings +import logging + +log = logging.getLogger('run_op') warnings.filterwarnings('ignore') @@ -83,6 +86,7 @@ class AutoTransformers(NNOperator): os.makedirs(path, exist_ok=True) name = self.model_name.replace('/', '-') path = os.path.join(path, name) + inputs = self.tokenizer('[CLS]', return_tensors='pt') # a dictionary if format == 'pytorch': path = path + '.pt' @@ -101,37 +105,39 @@ class AutoTransformers(NNOperator): raise RuntimeError(f'Fail to save as torchscript: {e}.') elif format == 'onnx': path = path + '.onnx' + input_names = list(inputs.keys()) + dynamic_axes = {} + for i_n in input_names: + dynamic_axes[i_n] = {0: 'batch_size', 1: 'sequence_length'} try: + output_names = ['last_hidden_state'] + for o_n in output_names: + dynamic_axes[o_n] = {0: 'batch_size', 1: 'sequence_length'} torch.onnx.export(self.model, tuple(inputs.values()), path, - input_names=list(inputs.keys()), - output_names=["last_hidden_state"], - dynamic_axes={ - "input_ids": {0: "batch_size", 1: "input_length"}, - "token_type_ids": {0: "batch_size", 1: "input_length"}, - "attention_mask": {0: "batch_size", 1: "input_length"}, - "last_hidden_state": {0: "batch_size"}, - }, - opset_version=13, + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + opset_version=11, do_constant_folding=True, # enable_onnx_checker=True, ) except Exception as e: print(e, '\nTrying with 2 outputs...') + output_names = ['last_hidden_state', 'pooler_output'] + for o_n in output_names: + dynamic_axes[o_n] = {0: 'batch_size', 1: 'sequence_length'} torch.onnx.export(self.model, tuple(inputs.values()), path, - input_names=["input_ids", "token_type_ids", "attention_mask"], # list(inputs.keys()) - output_names=["last_hidden_state", "pooler_output"], - opset_version=13, - dynamic_axes={ - "input_ids": {0: "batch_size", 1: "input_length"}, - "token_type_ids": {0: "batch_size", 1: "input_length"}, - "attention_mask": {0: "batch_size", 1: "input_length"}, - "last_hidden_state": {0: "batch_size"}, - "pooler_outputs": {0: "batch_size"} - }) + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + opset_version=11, + do_constant_folding=True, + # enable_onnx_checker=True, + ) # todo: elif format == 'tensorrt': else: log.error(f'Unsupported format "{format}".') diff --git a/test_onnx.py b/test_onnx.py index b586b03..239d2b5 100644 --- a/test_onnx.py +++ b/test_onnx.py @@ -13,9 +13,12 @@ import psutil # full_models = AutoTransformers.supported_model_names() # checked_models = AutoTransformers.supported_model_names(format='onnx') # models = [x for x in full_models if x not in checked_models] -models = ['distilbert-base-cased'] +models = ['bert-base-cased', 'distilbert-base-cased'] +test_txt = '[UNK]' atol = 1e-3 log_path = 'transformers_onnx.log' +f = open('onnx.csv', 'w+') +f.write('model,load_op,save_onnx,check_onnx,run_onnx,accuracy\n') logger = logging.getLogger('transformers_onnx') logger.setLevel(logging.DEBUG) @@ -30,24 +33,31 @@ ch.setFormatter(formatter) logger.addHandler(ch) logger.debug(f'machine: {platform.platform()}-{platform.processor()}') -logger.debug(f'free/available/total mem: {round(psutil.virtual_memory().free / (1024.0 **3))}' - f'/{round(psutil.virtual_memory().available / (1024.0 **3))}' - f'/{round(psutil.virtual_memory().total / (1024.0 **3))} GB') +logger.debug(f'free/available/total mem: {round(psutil.virtual_memory().free / (1024.0 ** 3))}' + f'/{round(psutil.virtual_memory().available / (1024.0 ** 3))}' + f'/{round(psutil.virtual_memory().total / (1024.0 ** 3))} GB') logger.debug(f'cpu: {psutil.cpu_count()}') + +status = None for name in models: logger.info(f'***{name}***') saved_name = name.replace('/', '-') + if status: + f.write(','.join(status) + '\n') + status = [name] + ['fail'] * 5 try: op = AutoTransformers(model_name=name) - out1 = op('hello, world.') + out1 = op(test_txt) logger.info('OP LOADED.') + status[1] = 'success' except Exception as e: logger.error(f'FAIL TO LOAD OP: {e}') continue try: op.save_model(format='onnx') logger.info('ONNX SAVED.') + status[2] = 'success' except Exception as e: logger.error(f'FAIL TO SAVE ONNX: {e}') continue @@ -59,21 +69,27 @@ for name in models: saved_onnx = onnx.load(f'saved/onnx/{saved_name}.onnx', load_external_data=False) onnx.checker.check_model(saved_onnx) logger.info('ONNX CHECKED.') + status[3] = 'success' except Exception as e: logger.error(f'FAIL TO CHECK ONNX: {e}') continue try: sess = onnxruntime.InferenceSession(f'saved/onnx/{saved_name}.onnx', providers=onnxruntime.get_available_providers()) - inputs = op.tokenizer('hello, world.', return_tensors='np') - out2 = sess.run(output_names=["last_hidden_state"], input_feed=dict(inputs)) + inputs = op.tokenizer(test_txt, return_tensors='np') + out2 = sess.run(output_names=['last_hidden_state'], input_feed=dict(inputs)) logger.info('ONNX WORKED.') + status[4] = 'success' if numpy.allclose(out1, out2, atol=atol): logger.info('Check accuracy: OK') + status[5] = 'success' else: logger.info(f'Check accuracy: atol is larger than {atol}.') except Exception as e: logger.error(f'FAIL TO RUN ONNX: {e}') continue +if status: + f.write(','.join(status) + '\n') + print('Finished.')