Add script to benchmark bert and mobilenet (#1594)

* Add script for bert and mobilenet

Signed-off-by: xavier dupré <[email protected]>

* add script to convert tfhub models

Signed-off-by: xavier dupré <[email protected]>

* add more scripts

Signed-off-by: xavier dupré <[email protected]>

* add all scripts

Signed-off-by: xavier dupré <[email protected]>

Co-authored-by: xavier dupré <[email protected]>
Co-authored-by: TomWildenhain-Microsoft <[email protected]>

Author: 3 people

.gitignore

@@ -13,3 +13,7 @@ __pycache__
 .eggs
 *.egg-info
 run.sh
+tests/tfhub/*/*.onnx
+tests/tfhub/*/*.tar.gz
+tests/tfhub/*/*.tflite
+tests/tfhub/*/**

tests/tfhub/_tools.py

@@ -11,21 +11,60 @@
 import zipfile
 import subprocess
 import datetime
+from collections import OrderedDict
 import numpy
 from tqdm import tqdm
 import onnxruntime
 
 
-def generate_random_images(shape=(1, 100, 100, 3), n=10, dtype=numpy.float32):
+def generate_random_images(shape=(1, 100, 100, 3), n=10, dtype=numpy.float32, scale=255):
     imgs = []
     for i in range(n):
         sh = shape
-        img = numpy.clip(numpy.abs(numpy.random.randn(*sh)), 0, 1) * 255
+        img = numpy.clip(numpy.abs(numpy.random.randn(*sh)), 0, 1) * scale
         img = img.astype(dtype)
         imgs.append(img)
     return imgs
 
 
+def generate_text_inputs():
+    """
+    preprocessor = hub.load("http://tfhub.dev/tensorflow/albert_en_preprocess/3")
+    encoder = hub.load("https://tfhub.dev/tensorflow/albert_en_xlarge/3")
+    sentences = tf.constant(["Hi I'm some text"])
+    embedded_inputs = {k: v.numpy() for k, v in preprocessor(sentences).items()}
+    """
+    one = OrderedDict([
+        ('input_word_ids', numpy.array([[
+            2, 4148, 31, 22, 79, 109, 1854, 3, 0, 0, 0, 
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0,0]]).reshape((1, -1))),
+        ('input_type_ids', numpy.array([[
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).reshape((1, -1))),
+        ('input_mask', numpy.array([[
+            1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).reshape((1, -1)))])
+    return [one for i in range(10)]
+
+
 def measure_time(fct, imgs, n=50, timeout=15):
     """
     Runs *n* times the same function taking one parameter
@@ -89,16 +128,22 @@ def download_tflite(url, dest, verbose=True):
     return fpath
 
 
-def convert_model(model_name, output_path, opset=13, verbose=True):
+def convert_model(model_name, output_path, opset=13, tag=None, verbose=True):
     """
     Converts the downloaded model into ONNX.
     """
+    ext = os.path.splitext(output_path)[-1]
+    large_model = ext == ".zip"
     if not os.path.exists(output_path):
         begin = datetime.datetime.now()
         cmdl = ['-m', 'tf2onnx.convert', '--saved-model',
                 '"%s"' % os.path.abspath(model_name).replace("\\", "/"),
                 '--output', '"%s"' % os.path.abspath(output_path).replace("\\", "/"),
                 '--opset', "%d" % opset]
+        if tag is not None:
+            cmdl.append('--tag="%s"' % tag)
+        if large_model:
+            cmdl.append('--large_model')
         if verbose:
             print("cmd: python %s" % " ".join(cmdl))
         pproc = subprocess.Popen(
@@ -151,7 +196,7 @@ def check_discrepencies(out1, out2, threshold=1e-3):
 
 
 def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
-              signature=None):
+              signature=None, tag=None, output_name=None, ort_name=None):
     """
     Runs a simple benchmark.
     Goes through every steps (download, convert).
@@ -164,10 +209,21 @@ def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
     # Converts the model.
     if verbose:
         print("Convert model in %r." % dest)
-    convert_model(tname, onnx_name, opset)
+    convert_model(tname, onnx_name, opset, tag=tag)
     if verbose:
         print("Created %r." % onnx_name)
 
+    # unzip large_model
+    ext = os.path.splitext(onnx_name)[-1]
+    if ext == ".zip":
+        onnx_name_unzipped = os.path.join(dest, "large_model", "__MODEL_PROTO.onnx")
+        if not os.path.exists(onnx_name_unzipped):
+            if verbose:
+                print("Unzip model in %r." % os.path.join(dest, "large_model"))
+            with zipfile.ZipFile(onnx_name, 'r') as z:
+              z.extractall(os.path.join(dest, "large_model"))
+        onnx_name = onnx_name_unzipped
+
     # Benchmarks both models.
     ort = onnxruntime.InferenceSession(onnx_name)
 
@@ -180,19 +236,37 @@ def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
             print("  {}: {}, {}".format(a.name, a.type, a.shape))
 
     # onnxruntime
-    input_name = ort.get_inputs()[0].name
-    fct_ort = lambda img: ort.run(None, {input_name: img})[0]
+    if output_name is None or ort_name is None:
+        index = 0
+    else:
+        output_names = [o.name for o in ort.get_outputs()]
+        if output_name in output_names:
+            index = output_names.index(output_name)
+        elif ort_name in output_names:
+            index = output_names.index(ort_name)
+        else:
+            index = 0
+    if isinstance(imgs[0], dict):
+        fct_ort = lambda img: ort.run(None, img)[index]
+    else:
+        input_name = ort.get_inputs()[0].name
+        fct_ort = lambda img: ort.run(None, {input_name: img})[index]
     results_ort, duration_ort = measure_time(fct_ort, imgs)
     if verbose:
         print("ORT", len(imgs), duration_ort)
 
     # tensorflow
     import tensorflow_hub as hub
     from tensorflow import convert_to_tensor
+    if isinstance(imgs[0], OrderedDict):
+        imgs_tf = [
+            OrderedDict((k, convert_to_tensor(v)) for k, v in img.items())
+            for img in imgs]
+    else:
+        imgs_tf = [convert_to_tensor(img) for img in imgs]
     model = hub.load(url.split("?")[0])
     if signature is not None:
-        model = model.signatures['serving_default']
-    imgs_tf = [convert_to_tensor(img) for img in imgs]
+        model = model.signatures[signature]
     results_tf, duration_tf = measure_time(model, imgs_tf)
 
     if verbose:
@@ -204,13 +278,27 @@ def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
     # checks discrepencies
     res = model(imgs_tf[0])
     if isinstance(res, dict):
-        if len(res) != 1:
-            raise NotImplementedError("TF output contains more than one output: %r." % res)
-        output_name = ort.get_outputs()[0].name
+        if output_name is None:
+            if len(res) != 1:
+                raise NotImplementedError(
+                    "TF output contains more than one output=%r and output names=%r." % (
+                        list(res), [o.name for o in ort.get_outputs()]))
+            else:
+                output_name = ort.get_outputs()[0].name
         if output_name not in res:
             raise AssertionError("Unable to find output %r in %r." % (output_name, list(sorted(res))))
         res = res[output_name]
-    check_discrepencies(fct_ort(imgs[0]), res.numpy(), threshold)
+    try:
+        check_discrepencies(fct_ort(imgs[0]), res.numpy(), threshold)
+    except AttributeError as e:
+        raise AssertionError(
+            "Unable to check discrepencies for res=%r." % res) from e
+    except AssertionError as e:
+        output_names = [o.name for o in ort.get_outputs()]
+        res = ort.run(None, imgs[0])
+        for i, r in enumerate(res):
+            print("ORT %d: %s: %r: %r" % (i, output_names[i], r.dtype, r.shape))
+        raise e
     return duration_ort, duration_tf
 
 
@@ -252,10 +340,15 @@ def benchmark_tflite(url, dest, onnx_name, opset, imgs, verbose=True, threshold=
     # tensorflow
     import tensorflow_hub as hub
     from tensorflow import convert_to_tensor
+    if isinstance(imgs[0], OrderedDict):
+        imgs_tf = [
+            OrderedDict((k, convert_to_tensor(v)) for k, v in img.items())
+            for img in imgs]
+    else:
+        imgs_tf = [convert_to_tensor(img) for img in imgs]
     model = hub.load(url.split("?")[0])
     if signature is not None:
         model = model.signatures['serving_default']
-    imgs_tf = [convert_to_tensor(img) for img in imgs]
     results_tf, duration_tf = measure_time(model, imgs_tf)
 
     if verbose:

tests/tfhub/tfhub_albert_en_xlarge.py

@@ -1,19 +1,27 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import numpy
-from _tools import generate_random_images, benchmark
+import numpy.random as rnd
+from collections import OrderedDict
+from _tools import generate_text_inputs, benchmark
 
 
 def main(opset=13):
     url = "https://tfhub.dev/tensorflow/albert_en_xlarge/3?tf-hub-format=compressed"
     dest = "tf-albert-en-xlarge"
     name = "albert-en-xlarge"
-    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
+    onnx_name = os.path.join(dest, "%s-%d.zip" % (name, opset))
 
-    imgs = generate_random_images(shape=(1, 256, 256, 3), dtype=numpy.int32)
+    inputs = generate_text_inputs()
+    benchmark(url, dest, onnx_name, opset, inputs, output_name="pooled_output")
 
-    benchmark(url, dest, onnx_name, opset, imgs,
-              signature='serving_default')
+    inputs = [OrderedDict([
+        ('input_word_ids', numpy.array([rnd.randint(0, 1000) for i in range(0, 128)], dtype=numpy.int32).reshape((1, -1))),
+        ('input_mask', numpy.array([rnd.randint(0, 1) for i in range(0, 128)], dtype=numpy.int32).reshape((1, -1))),
+        ('input_type_ids', numpy.array([i//5 for i in range(0, 128)], dtype=numpy.int32).reshape((1, -1)))
+    ]) for i in range(0, 10)]
+   
+    benchmark(url, dest, onnx_name, opset, inputs, output_name="pooled_output")
 
 
 if __name__ == "__main__":

tests/tfhub/tfhub_albert_en_xlarge_keras.py

@@ -0,0 +1,54 @@
+# Adapted the sample code on https://tfhub.dev/tensorflow/albert_en_xlarge/3
+import tensorflow_text as text
+import tensorflow as tf
+import tensorflow_hub as hub
+
+# Using hub.load instead of KerasLayer lets us easily intercept the results of the
+# preprocessor before passing it to the encoder
+preprocessor = hub.load("http://tfhub.dev/tensorflow/albert_en_preprocess/3")
+encoder = hub.load("https://tfhub.dev/tensorflow/albert_en_xlarge/3")
+sentences = tf.constant(["Hi I'm some text"])
+
+embedded_inputs = {k: v.numpy() for k, v in preprocessor(sentences).items()}
+print("Inputs")
+print(embedded_inputs)
+expected_output = encoder(embedded_inputs)["pooled_output"].numpy()
+
+# Now make an actual keras layer for the part we want to convert
+encoder = hub.KerasLayer(
+    "https://tfhub.dev/tensorflow/albert_en_xlarge/3",
+    trainable=True)
+
+# To convert it to a model, we need the input shapes/types. These can be
+# determined from the types/shapes/names of embedded_inputs. Remove the batch dim from the shapes.
+encoder_inputs = {
+    "input_word_ids": tf.keras.Input(shape=[None], dtype="int32", name="input_word_ids"),
+    "input_mask": tf.keras.Input(shape=[None], dtype="int32", name="input_mask"),
+    "input_type_ids": tf.keras.Input(shape=[None], dtype="int32", name="input_type_ids"),
+}
+encoder_outputs = encoder(encoder_inputs)["pooled_output"]
+encoding_model = tf.keras.Model(encoder_inputs, encoder_outputs)
+
+
+import tf2onnx
+import onnxruntime as ort
+import zipfile
+import os
+print("Converting")
+
+dest = "tf-albert-en-xlarge"
+if not os.path.exists(dest):
+    os.makedirs(dest)
+dest_name = os.path.join(dest, "albert_en_xlarge.zip")
+
+# This model is a large model.
+tf2onnx.convert.from_keras(encoding_model, opset=13, large_model=True, output_path=dest_name)
+# To run the model in ORT we need to unzip it.
+with zipfile.ZipFile(dest_name, 'r') as z:
+  z.extractall(os.path.join(dest, "albert_en_xlarge"))
+sess = ort.InferenceSession(os.path.join(dest, "albert_en_xlarge", "__MODEL_PROTO.onnx"))
+ort_output = sess.run(None, embedded_inputs)
+print("Actual")
+print(ort_output[0])
+print("Expected")
+print(expected_output)

tests/tfhub/tfhub_bert_en_wwm_uncased.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+from collections import OrderedDict
+import numpy
+import numpy.random as rnd
+from _tools import generate_random_images, benchmark
+
+
+def main(opset=13):
+    url = "https://tfhub.dev/tensorflow/bert_en_wwm_uncased_L-24_H-1024_A-16/4?tf-hub-format=compressed"
+    dest = "tf-bert-en-wwm-uncased-L-24-H-1024-A-16"
+    name = "bert-en-wwm-uncased-L-24-H-1024-A-16"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
+
+    inputs = [OrderedDict([
+        ('input_word_ids', numpy.array([rnd.randint(0, 1000) for i in range(0, 32)], dtype=numpy.int32).reshape((1, -1))),
+        ('input_mask', numpy.array([rnd.randint(0, 1) for i in range(0, 32)], dtype=numpy.int32).reshape((1, -1))),
+        ('input_type_ids', numpy.array([i//5 for i in range(0, 32)], dtype=numpy.int32).reshape((1, -1)))
+    ]) for i in range(0, 10)]
+
+    benchmark(url, dest, onnx_name, opset, inputs, output_name="pooled_output")
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_blazeposedetector.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import numpy
+from _tools import generate_random_images, benchmark
+
+
+def main(opset=13):
+    url = "https://tfhub.dev/mediapipe/tfjs-model/blazeposedetector/1/default/1?tfjs-format=compressed"
+    dest = "tf-blazeposedetector"
+    name = "blazeposedetector"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
+
+    imgs = generate_random_images(shape=(1, 513, 513, 3), scale=1.)
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_enformer.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import numpy
+from _tools import generate_random_images, benchmark
+
+
+def main(opset=13):
+    url = "https://tfhub.dev/deepmind/enformer/1?tf-hub-format=compressed"
+    dest = "tf-enformer"
+    name = "enformer"
+    onnx_name = os.path.join(dest, "%s-%d.zip" % (name, opset))
+
+    imgs = generate_random_images(shape=(1, 224, 224, 3))
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_humpback_whale.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import numpy
+from _tools import generate_random_images, benchmark
+
+
+def main(opset=13):
+    url = "https://tfhub.dev/google/humpback_whale/1?tf-hub-format=compressed"
+    dest = "tf-humpback-whale"
+    name = "humpback-whale"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
+
+    imgs = generate_random_images(shape=(1, 331, 331, 3))
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()