Add script for bert and mobilenet

Author: sdpython

.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,16 +11,17 @@
 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
@@ -180,19 +181,27 @@ 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 isinstance(imgs[0], dict):
+        fct_ort = lambda img: ort.run(None, img)[0]
+    else:
+        input_name = ort.get_inputs()[0].name
+        fct_ort = lambda img: ort.run(None, {input_name: img})[0]
     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]
     results_tf, duration_tf = measure_time(model, imgs_tf)
 
     if verbose:
@@ -205,7 +214,9 @@ def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
     res = model(imgs_tf[0])
     if isinstance(res, dict):
         if len(res) != 1:
-            raise NotImplementedError("TF output contains more than one output: %r." % res)
+            raise NotImplementedError(
+                "TF output contains more than one output=%r and output names=%r." % (
+                    res, [o.name for o in ort.get_outputs()]))
         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))))
@@ -252,10 +263,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_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)
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_mobilenet_v3_large_075_224.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/imagenet/mobilenet_v3_large_075_224/classification/5?tf-hub-format=compressed"
+    dest = "tf-mobilenet-v3-large-075-224"
+    name = "mobilenet-v3-large-075-224"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
+
+    imgs = generate_random_images(shape=(1, 224, 224, 3), scale=1.)
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()