Add ONNX Runtime, statistics visualization

Author: mtszkw

.gitignore

@@ -1,2 +1,3 @@
 *.lock
-*.pt
+*.pt
+*.onnx

Pipfile

@@ -7,6 +7,10 @@ name = "pypi"
 transformers = {extras = ["torch"], version = "*"}
 scipy = "*"
 matplotlib = "*"
+seaborn = "*"
+pandas = "*"
+onnxruntime = "*"
+onnx = "*"
 
 [dev-packages]
 

inference.py

@@ -0,0 +1,43 @@
+import numpy as np
+# https://stackoverflow.com/questions/7370801/how-to-measure-elapsed-time-in-python
+from timeit import default_timer as timer
+
+
+class PythonInference:
+    def __init__(self, model):
+        self.model = model
+
+    def check_inference_time(self, tokenized_input: dict):
+        t = timer()
+        output = self.model(**tokenized_input)
+        elapsed_time = timer()-t
+        return elapsed_time
+
+    def check_inference_time_all(self, tokenized_inputs: list, num_experiments=1):
+        time_measurements = []
+        for i in range(num_experiments):
+            time_measurements.append([self.check_inference_time(x) for x in tokenized_inputs])
+        return np.array(time_measurements)
+        # t = timer()
+        # output = self.model(**tokenized_input)
+        # elapsed_time = timer()-t
+        # return elapsed_time
+
+
+class OnxxInference:
+    def __init__(self, session):
+        self.session = session
+
+    def check_inference_time(self, tokenized_input):
+        t = timer()
+        output = self.session.run(None, tokenized_input)
+        elapsed_time = timer()-t
+        return elapsed_time
+    
+    def check_inference_time_all(self, tokenized_inputs: list, num_experiments=1):
+        time_measurements = []
+        for i in range(num_experiments):
+            time_measurements.append(
+                [self.check_inference_time({name: np.atleast_2d(value) for name, value in x.items()}) for x in tokenized_inputs])
+        return np.array(time_measurements)
+        

run_sentiment_classifier.py

@@ -1,88 +1,71 @@
-import numpy as np
 import matplotlib.pyplot as plt
-# https://stackoverflow.com/questions/7370801/how-to-measure-elapsed-time-in-python
-from timeit import default_timer as timer
+import numpy as np
+import pandas as pd
+import seaborn as sns
 
 import torch
 from transformers import AutoTokenizer
 from transformers import AutoModelForSequenceClassification
 
-from utils import preprocess, download_label_mapping, output_vector_to_labels
+from onnxruntime import ExecutionMode, InferenceSession, SessionOptions
 
+from inference import PythonInference, OnxxInference
+from utils import preprocess, download_label_mapping, output_vector_to_labels, measurements_to_dataframe
 
 def read_test_sequences(path: str):
     with open(path, 'r') as f:
         sequences = [x.rstrip() for x in f.readlines()]
         return sequences
 
-
-def run_model(model, tokenized_input):
-    output = model(**tokenized_input)
-    return output_vector_to_labels(output, download_label_mapping())
-
-
-def check_inference_time(model, tokenized_input):
-    t = timer()
-    scores = run_model(model, tokenized_input)
-    elapsed_time = timer()-t
-    return elapsed_time
-
-
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment", torchscript=True)
     clf = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment", torchscript=True)
 
-    n_experiments = 5
     input_texts = [preprocess(x) for x in read_test_sequences("test_sequences.txt")]
-
-    # 1. Eager
-    eager_measurements = np.zeros((n_experiments, len(input_texts)))
     tokenized_inputs = [tokenizer(x, return_tensors='pt') for x in input_texts]
+    
+    n_experiments = 2
+    indices = np.tile(np.arange(len(input_texts)), n_experiments)
 
-    for i in range(n_experiments):
-        # outputs = [run_model(clf, x) for x in tokenized_inputs]
-        eager_measurements[i] = [check_inference_time(clf, x) for x in tokenized_inputs]
-        # for inp, out in zip(input_texts, outputs):
-            # print(inp, '\n', out, '\n')
-        # print(output_times)
+    
+    # 1. Eager
+    eager_model = PythonInference(clf)
+    eager_measurements = eager_model.check_inference_time_all(tokenized_inputs, n_experiments)
+    df_eager = measurements_to_dataframe(eager_measurements.flatten(), indices)
+    df_eager['Mode'] = 'Eager'
 
 
     # 2. TorchScript (JIT)
-    script_measurements = np.zeros((n_experiments, len(input_texts)))
-    tokenized_inputs = [tokenizer(x, return_tensors='pt') for x in input_texts]
-    traced_model = torch.jit.trace(clf, (tokenized_inputs[0]['input_ids'], tokenized_inputs[0]['attention_mask']))
-    # torch.jit.save(traced_model, "traced_twitter_roberta_base_sentiment.pt")
-    # loaded_model = torch.jit.load("traced_twitter_roberta_base_sentiment.pt")
-
-    for i in range(n_experiments):
-        # outputs = [run_model(traced_model, x) for x in tokenized_inputs]
-        script_measurements[i] = [check_inference_time(traced_model, x) for x in tokenized_inputs]
-        # for inp, out in zip(input_texts, outputs):
-            # print(inp, '\n', out, '\n')
-        # print(output_times)
+    traced_model = PythonInference(model=torch.jit.trace(clf, (tokenized_inputs[0]['input_ids'], tokenized_inputs[0]['attention_mask'])))
+    script_measurements = traced_model.check_inference_time_all(tokenized_inputs, n_experiments)
+    df_script = measurements_to_dataframe(script_measurements.flatten(), indices)
+    df_script['Mode'] = 'Script'
 
-    print(eager_measurements)
-    print(script_measurements)
-
-    # Box Plot
 
-    eager_avgs = np.mean(eager_measurements, axis=0)
-    script_avgs = np.mean(script_measurements, axis=0)
-    print(eager_avgs)
-    print(script_avgs)
-
-    # Scatter Plot
+    # 3. ONNX Runtime
+    model = OnxxInference(session=InferenceSession("onnx_model/twitter-roberta-base-sentiment-optimized-quantized.onnx"))
+    onnx_measurements = model.check_inference_time_all(tokenized_inputs, n_experiments)
+    df_onnx = measurements_to_dataframe(onnx_measurements.flatten(), indices)
+    df_onnx['Mode'] = 'ONNX'
 
-    indices = np.tile(np.arange(len(input_texts)), n_experiments)
-    eager_measurements = eager_measurements.flatten()
-    script_measurements = script_measurements.flatten()
-    print(indices)
-    print(eager_measurements)
 
+    # Statistics
     plt.style.use('seaborn')
-    plt.scatter(indices, eager_measurements, label='Eager mode')
-    plt.scatter(indices, script_measurements, label='Script mode')
+    plt.figure()
+    plt.scatter(x=df_eager['SequenceId'], y=df_eager['TimeInSeconds'], label='Eager mode')
+    plt.scatter(x=df_script['SequenceId'], y=df_script['TimeInSeconds'], label='Script mode')
+    plt.scatter(x=df_onnx['SequenceId'], y=df_onnx['TimeInSeconds'], label='ONNX mode')
     plt.xlabel('Sequence ID')
     plt.ylabel('Inference time [s]')
     plt.legend()
     plt.show()
+
+
+    plt.figure()
+    df_all = pd.concat([df_eager, df_script, df_onnx])
+    df_all.groupby('Mode').mean().TimeInSeconds.plot(kind='bar')
+    plt.title('Avg. inference time in seconds')
+    plt.ylabel('Inference time [s]')
+    plt.show()
+
+    # Box plots

utils.py

@@ -1,5 +1,6 @@
 import csv
 import numpy as np
+import pandas as pd
 import urllib.request
 from scipy.special import softmax
 
@@ -22,7 +23,10 @@ def download_label_mapping():
     return labels
 
 def output_vector_to_labels(output, labels_map):
-    scores = output[0][0].detach().numpy()
+    if isinstance(output[0][0], np.ndarray):
+        scores = output[0][0]
+    else:
+        scores = output[0][0].detach().numpy()
     scores = softmax(scores)
 
     scores_map = {}
@@ -33,4 +37,12 @@ def output_vector_to_labels(output, labels_map):
         s = scores[ranking[i]]
         scores_map[l] = np.round(float(s), 4)
         # print(f"{i+1}) {l} {np.round(float(s), 4)}")
-    return scores_map
+    return scores_map
+
+def measurements_to_dataframe(measurements, indices):
+    data = []
+    for seq_id, time in zip(indices, measurements):
+        data.append([seq_id, time, 'Unknown'])
+    df = pd.DataFrame(data, columns=['SequenceId', 'TimeInSeconds', 'Mode'])
+    print(df)
+    return df