Merge pull request #1587 from xadupre/bench2

Add yamnet to the list of tfhub tested models

Author: xadupre

tests/tfhub/_tools.py

@@ -73,6 +73,22 @@ def download_model(url, dest, verbose=True):
     return fpath, tname
 
 
+def download_tflite(url, dest, verbose=True):
+    """
+    Downloads a model from tfhub.
+    The function assumes the format is `.tflite`.
+    """
+    if not os.path.exists(dest):
+        os.makedirs(dest)
+    fpath = os.path.join(dest, "model.tflite")
+    if not os.path.exists(fpath):
+        from tf2onnx import utils
+        if verbose:
+            print("Download %r." % fpath)
+        utils.get_url(url, fpath)
+    return fpath
+
+
 def convert_model(model_name, output_path, opset=13, verbose=True):
     """
     Converts the downloaded model into ONNX.
@@ -97,6 +113,30 @@ def convert_model(model_name, output_path, opset=13, verbose=True):
             print("Duration %r." % (datetime.datetime.now() - begin))
 
 
+def convert_tflite(model_name, output_path, opset=13, verbose=True):
+    """
+    Converts the downloaded model into ONNX.
+    """
+    if not os.path.exists(output_path):
+        begin = datetime.datetime.now()
+        cmdl = ['-m', 'tf2onnx.convert', '--tflite',
+                '"%s"' % os.path.abspath(model_name).replace("\\", "/"),
+                '--output', '"%s"' % os.path.abspath(output_path).replace("\\", "/"),
+                '--opset', "%d" % opset]
+        if verbose:
+            print("cmd: python %s" % " ".join(cmdl))
+        pproc = subprocess.Popen(
+            cmdl, shell=True, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+            executable=sys.executable.replace("pythonw", "python"))
+        stdoutdata, stderrdata = pproc.communicate()
+        if verbose:
+            print('--OUT--')
+            print(stdoutdata.decode('ascii'))
+            print('--ERR--')
+            print(stderrdata.decode('ascii'))
+            print("Duration %r." % (datetime.datetime.now() - begin))
+
+
 def check_discrepencies(out1, out2, threshold=1e-3):
     """
     Compares two tensors. Raises an exception if it fails.
@@ -172,3 +212,66 @@ def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
         res = res[output_name]
     check_discrepencies(fct_ort(imgs[0]), res.numpy(), threshold)
     return duration_ort, duration_tf
+
+
+def benchmark_tflite(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3):
+    """
+    Runs a simple benchmark with a tflite model.
+    Goes through every steps (download, convert).
+    Skips them if already done.
+    """
+    tname = download_tflite(url, dest)
+    if verbose:
+        print("Created %r." % tname)
+
+    # Converts the model.
+    if verbose:
+        print("Convert model in %r." % dest)
+    convert_tflite(tname, onnx_name, opset)
+    if verbose:
+        print("Created %r." % onnx_name)
+
+    # Benchmarks both models.
+    ort = onnxruntime.InferenceSession(onnx_name)
+
+    if verbose:
+        print("ONNX inputs:")
+        for a in ort.get_inputs():
+            print("  {}: {}, {}".format(a.name, a.type, a.shape))
+        print("ONNX outputs:")
+        for a in ort.get_outputs():
+            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]
+    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
+    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:
+        print("TF", len(imgs), duration_tf)
+        mean_ort = sum(duration_ort) / len(duration_ort)
+        mean_tf = sum(duration_tf) / len(duration_tf)
+        print("ratio ORT=%r / TF=%r = %r" % (mean_ort, mean_tf, mean_ort / mean_tf))
+
+    # 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 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)
+    return duration_ort, duration_tf

tests/tfhub/tfhub_albert_en_xlarge.py

@@ -0,0 +1,20 @@
+# 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/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))
+
+    imgs = generate_random_images(shape=(1, 256, 256, 3), dtype=numpy.int32)
+
+    benchmark(url, dest, onnx_name, opset, imgs,
+              signature='serving_default')
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_esrgan.py

@@ -1,14 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import numpy
-from _tools import generate_random_images, measure_time, download_model, convert_model, benchmark
+from _tools import generate_random_images, benchmark
 
-url = "https://tfhub.dev/captain-pool/esrgan-tf2/1?tf-hub-format=compressed"
-dest = "tf-esrgan-tf2"
-name = "esrgan-tf2"
-opset = 13
-onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-imgs = generate_random_images()
+def main(opset=13):
+    url = "https://tfhub.dev/captain-pool/esrgan-tf2/1?tf-hub-format=compressed"
+    dest = "tf-esrgan-tf2"
+    name = "esrgan-tf2"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-benchmark(url, dest, onnx_name, opset, imgs)
+    imgs = generate_random_images()
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_resnet_v2_152.py

@@ -1,14 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import numpy
-from _tools import generate_random_images, measure_time, download_model, convert_model, benchmark
+from _tools import generate_random_images, benchmark
 
-url = "https://tfhub.dev/google/imagenet/resnet_v2_152/classification/5?tf-hub-format=compressed"
-dest = "tf-resnet_v2_152"
-name = "resnet_v2_152"
-opset = 13
-onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-imgs = generate_random_images(shape=(1, 224, 224, 3))
+def main(opset=13):
+    url = "https://tfhub.dev/google/imagenet/resnet_v2_152/classification/5?tf-hub-format=compressed"
+    dest = "tf-resnet_v2_152"
+    name = "resnet_v2_152"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-benchmark(url, dest, onnx_name, opset, imgs)
+    imgs = generate_random_images(shape=(1, 224, 224, 3))
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_spam_detection.py

@@ -2,14 +2,19 @@
 import os
 import random
 import numpy
-from _tools import generate_random_images, measure_time, download_model, convert_model, benchmark
+from _tools import generate_random_images, benchmark
 
-url = "https://tfhub.dev/tensorflow/tutorials/spam-detection/1?tf-hub-format=compressed"
-dest = "tf-spam-detection"
-name = "spam-detection"
-opset = 13
-onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-imgs = generate_random_images((1, 20), dtype=numpy.int32)
+def main(opset=13):
+    url = "https://tfhub.dev/tensorflow/tutorials/spam-detection/1?tf-hub-format=compressed"
+    dest = "tf-spam-detection"
+    name = "spam-detection"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-benchmark(url, dest, onnx_name, opset, imgs)
+    imgs = generate_random_images((1, 20), dtype=numpy.int32)
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_thunder.py

@@ -1,15 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import numpy
-from _tools import generate_random_images, measure_time, download_model, convert_model, benchmark
+from _tools import generate_random_images, benchmark
 
-url = "https://tfhub.dev/google/movenet/singlepose/thunder/3?tf-hub-format=compressed"
-dest = "tf-thunder"
-name = "thunder"
-opset = 13
-onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-imgs = generate_random_images(shape=(1, 256, 256, 3), dtype=numpy.int32)
+def main(opset=13):
+    url = "https://tfhub.dev/google/movenet/singlepose/thunder/3?tf-hub-format=compressed"
+    dest = "tf-thunder"
+    name = "thunder"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-benchmark(url, dest, onnx_name, opset, imgs,
-          signature='serving_default')
+    imgs = generate_random_images(shape=(1, 256, 256, 3), dtype=numpy.int32)
+
+    benchmark(url, dest, onnx_name, opset, imgs,
+              signature='serving_default')
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_yamnet.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import numpy
+from _tools import generate_random_images, benchmark_tflite
+
+
+def main(opset=13):
+    url = "https://tfhub.dev/google/coral-model/yamnet/classification/coral/1?coral-format=tflite"
+    dest = "tf-yamnet"
+    name = "yamnet"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
+
+    imgs = generate_random_images(shape=(1, 256, 256, 3), dtype=numpy.int32)
+
+    benchmark_tflite(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()