Fix log

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>

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}".')

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.')