[WIP] Fix tfhub scripts (enformer, ...) (#1637)

* Scripts and errors

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

* yamnet

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

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

Author: xadupre

tests/tfhub/_tools.py

@@ -128,27 +128,29 @@ def download_tflite(url, dest, verbose=True):
     return fpath
 
 
-def convert_model(model_name, output_path, opset=13, tag=None, verbose=True):
+def convert_model(model_name, output_path, opset=13, tag=None, signature=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("\\", "/"),
+        cmdl = ['python', '-m', 'tf2onnx.convert', '--saved-model',
+                '%s' % os.path.abspath(model_name).replace("\\", "/"),
+                '--output', '%s' % os.path.abspath(output_path).replace("\\", "/"),
                 '--opset', "%d" % opset]
+        if signature is not None:
+            cmdl.append('--signature_def=%s' % signature)
         if tag is not None:
-            cmdl.append('--tag="%s"' % tag)
+            cmdl.append('--tag=%s' % tag)
         if large_model:
             cmdl.append('--large_model')
         if verbose:
-            print("cmd: python %s" % " ".join(cmdl))
+            print("cmd: %s" % " ".join(cmdl))
         pproc = subprocess.Popen(
-            cmdl, shell=True, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
-            executable=sys.executable.replace("pythonw", "python"))
+            cmdl, shell=False, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+            executable=None)
         stdoutdata, stderrdata = pproc.communicate()
         if verbose:
             print('--OUT--')
@@ -164,15 +166,15 @@ def convert_tflite(model_name, output_path, opset=13, verbose=True):
     """
     if not os.path.exists(output_path):
         begin = datetime.datetime.now()
-        cmdl = ['-m', 'tf2onnx.convert', '--tflite',
+        cmdl = ['python', '-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))
+            print("cmd: %s" % " ".join(cmdl))
         pproc = subprocess.Popen(
-            cmdl, shell=True, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
-            executable=sys.executable.replace("pythonw", "python"))
+            cmdl, shell=False, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+            executable=None)
         stdoutdata, stderrdata = pproc.communicate()
         if verbose:
             print('--OUT--')
@@ -186,6 +188,20 @@ def check_discrepencies(out1, out2, threshold=1e-3):
     """
     Compares two tensors. Raises an exception if it fails.
     """
+    if isinstance(out1, list):
+        if len(out1) > 1:
+            if len(out1) != len(out2):
+                raise AssertionError(
+                    "Mismatched number of outputs, %d for ONNX, %d for TF." % (
+                        len(out1), len(out2)))
+            for i, (a, b) in enumerate(zip(out1, out2)):
+                try:
+                    check_discrepencies(out1[i], out2[i].numpy(), threshold=1e-3)
+                except AssertionError as e:
+                    raise AssertionError("Discrepency with output %d." % i) from e
+            return
+        else:
+            out1 = out1[0]
     if out1.dtype != out2.dtype:
         raise AssertionError("Type mismatch %r != %r." % (out1.dtype, out2.dtype))
     if out1.shape != out2.shape:
@@ -210,7 +226,7 @@ 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, tag=tag)
+    convert_model(tname, onnx_name, opset, tag=tag, signature=signature)
     if verbose:
         print("Created %r." % onnx_name)
 
@@ -254,9 +270,11 @@ def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
             index = 0
     if isinstance(imgs[0], dict):
         fct_ort = lambda img: ort.run(None, img)[index]
+        fct_orts = lambda img: ort.run(None, img)
     else:
         input_name = ort.get_inputs()[0].name
         fct_ort = lambda img: ort.run(None, {input_name: img})[index]
+        fct_orts = lambda img: ort.run(None, {input_name: img})
     results_ort, duration_ort = measure_time(fct_ort, imgs)
     if verbose:
         print("ORT", len(imgs), duration_ort)
@@ -294,8 +312,12 @@ def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
         if output_name not in res:
             raise AssertionError("Unable to find output %r in %r." % (output_name, list(sorted(res))))
         res = res[output_name]
+    res_ort = fct_orts(imgs[0])
     try:
-        check_discrepencies(fct_ort(imgs[0]), res.numpy(), threshold)
+        if len(res_ort) > 1:
+            check_discrepencies(res_ort, res, threshold)
+        else:
+            check_discrepencies(res_ort, res.numpy(), threshold)
     except AttributeError as e:
         raise AssertionError(
             "Unable to check discrepencies for res=%r." % res) from e
@@ -373,4 +395,4 @@ def benchmark_tflite(url, dest, onnx_name, opset, imgs, verbose=True, threshold=
             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
+    return duration_ort, duration_tf

tests/tfhub/tfhub_enformer.py

@@ -1,19 +1,60 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import numpy
-from _tools import generate_random_images, benchmark
+from numpy.testing import assert_almost_equal
+from onnxruntime import InferenceSession
+from _tools import generate_random_images, benchmark, measure_time
+from tensorflow import convert_to_tensor
+import tensorflow as tf
+import tensorflow_hub as hub
+import tf2onnx
 
 
 def main(opset=13):
+    print('[begin]')
     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))
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-    imgs = generate_random_images(shape=(1, 224, 224, 3))
+    model = None
+    if not os.path.exists(onnx_name):
+        if model is None:
+            model = hub.load("https://tfhub.dev/deepmind/enformer/1").model
 
-    benchmark(url, dest, onnx_name, opset, imgs)
+        tf2onnx.convert.from_function(
+            model.predict_on_batch,
+            [tf.TensorSpec([None, 393216, 4], tf.float32)],
+            opset=13, output_path=onnx_name)
 
+    # benchmark(url, dest, onnx_name, opset, imgs)
+    print("[generate dummy images]")
+    imgs = generate_random_images(shape=(1, 393216, 4), scale=0.)
+
+    ort = InferenceSession(onnx_name)
+    fct_ort = lambda img: ort.run(None, {'args_0': img})[0]
+
+    if model is None:
+        model = hub.load("https://tfhub.dev/deepmind/enformer/1").model
+
+    fct_tf = lambda img: model.predict_on_batch(img)
+
+    print('[benchmark tf]')
+    imgs_tf = [convert_to_tensor(img) for img in imgs]
+    results_tf, duration_tf = measure_time(fct_tf, imgs)
+    print("TF", len(imgs), duration_tf)
+
+    print('[benchmark ort]')
+    results_ort, duration_ort = measure_time(fct_ort, imgs)
+    print("ORT", len(imgs), duration_ort)    
+
+    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))
+
+    # discrepencies
+    assert_almost_equal(results_tf[0]['human'], results_ort[0], decimal=4)
+    print('[end]')
 
 if __name__ == "__main__":
     main()

tests/tfhub/tfhub_humpback_whale.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import numpy
+from onnxruntime import InferenceSession
 from _tools import generate_random_images, benchmark
 
 
@@ -10,12 +11,70 @@ def main(opset=13):
     name = "humpback-whale"
     onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
 
-    imgs = generate_random_images(shape=(1, 1024, 1))
-    inputs = [dict(waveform=img,
-                   context_step_samples=numpy.array(512, dtype=numpy.int64))
-              for img in imgs]
+    kind = "function"
+    if kind == "function":
+        import tensorflow as tf
+        import tensorflow_hub as hub
+        import tf2onnx
+        model = hub.load('https://tfhub.dev/google/humpback_whale/1')
+        FILENAME = 'gs://bioacoustics-www1/sounds/Cross_02_060203_071428.d20_7.wav'
+        waveform, sample_rate = tf.audio.decode_wav(tf.io.read_file(FILENAME))
+        waveform = tf.expand_dims(waveform, 0)  # makes a batch of size 1
+        context_step_samples = tf.cast(sample_rate, tf.int64)
+        print(waveform.dtype, waveform.shape, sample_rate.dtype, sample_rate.shape, sample_rate)
+        
+        spec = (tf.TensorSpec((None, ) + waveform.shape[-2:], tf.float32, name="waveform"),
+                tf.TensorSpec((1, 1), tf.int64, name="context_step_samples"))
+        inputs = {'waveform': waveform.numpy(),
+                  'context_step_samples': context_step_samples.numpy()}
+                
+        tf2onnx.convert.from_function(
+            model.signatures['score'], input_signature=spec, opset=13, output_path=onnx_name)
+        # AttributeError: '_WrapperFunction' object has no attribute 'get_concrete_function'
 
-    benchmark(url, dest, onnx_name, opset, inputs, optimize=False)
+        sess = InferenceSession(onnx_name)
+        got = sess.run(None, inputs)
+        print(got)
+        
+        score_fn = model.signatures['score']
+        scores = score_fn(waveform=waveform, context_step_samples=context_step_samples)
+    
+    if kind == "keras":
+        import tensorflow as tf
+        import tensorflow_hub as hub
+        import tf2onnx
+        model = hub.load('https://tfhub.dev/google/humpback_whale/1').model
+        FILENAME = 'gs://bioacoustics-www1/sounds/Cross_02_060203_071428.d20_7.wav'
+        waveform, sample_rate = tf.audio.decode_wav(tf.io.read_file(FILENAME))
+        waveform = tf.expand_dims(waveform, 0)  # makes a batch of size 1
+        context_step_samples = tf.cast(sample_rate, tf.int64)
+        print(waveform.dtype, waveform.shape, sample_rate.dtype, sample_rate.shape, sample_rate)
+
+        spec = (tf.TensorSpec((None, ) + waveform.shape[-2:], tf.float32, name="waveform"),
+                tf.TensorSpec((1, 1), tf.int64, name="context_step_samples"))
+        inputs = {'waveform': waveform.numpy(),
+                  'context_step_samples': context_step_samples.numpy()}
+                
+        tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path=onnx_name)
+        # AttributeError: '_UserObject' object has no attribute 'output_names'
+
+        sess = InferenceSession(onnx_name)
+        got = sess.run(None, inputs)
+        print(got)
+        
+        score_fn = model.signatures['score']
+        scores = score_fn(waveform=waveform, context_step_samples=context_step_samples)
+
+    if kind == 'cmd':
+        imgs = generate_random_images(shape=(1, 10000, 1), scale=1.)
+        inputs = [dict(waveform=img,
+                       context_step_samples=numpy.array(512, dtype=numpy.int64))
+                  for img in imgs]
+        benchmark(url, dest, onnx_name, opset, inputs, optimize=False,
+                  signature='score')
+        # onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: 
+        # [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'StatefulPartitionedCall/Reshape_1' Status Message: C:\xadupre\microsoft_xadupre\onnxruntime\onnxruntime\core\providers\cpu\tensor\reshape_helper.h:42 onnxruntime::ReshapeHelper::ReshapeHelper gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. 
+        # Input shape:{0,1}, requested shape:{1,1,1}
 
 
 if __name__ == "__main__":

tests/tfhub/tfhub_mobile_food_segmenter_V1.py

@@ -2,6 +2,8 @@
 import os
 import numpy
 from _tools import generate_random_images, benchmark
+import tf2onnx
+import onnxruntime as ort
 
 
 def main(opset=13):
@@ -12,7 +14,29 @@ def main(opset=13):
 
     imgs = generate_random_images(shape=(1, 513, 513, 3), scale=1.)
 
-    benchmark(url, dest, onnx_name, opset, imgs, tag='')
+    if True:
+        benchmark(url, dest, onnx_name, opset, imgs, tag='')
+        # The conversion works but tensorflow fails with
+        # TypeError: 'AutoTrackable' object is not callable
+
+    if True:
+        import tensorflow.compat.v2 as tf
+        import tensorflow_hub as hub
+
+        m = hub.KerasLayer('https://tfhub.dev/google/seefood/segmenter/mobile_food_segmenter_V1/1')
+        inputs = {
+            "X": tf.keras.Input(shape=[1, 513, 513, 3], dtype="float32", name="X"),
+        }
+        outputs = m(inputs)["default"]
+        # TypeError: pruned(images) missing required arguments: images
+        print(outputs)
+        model = tf.keras.Model(inputs, outputs)
+
+        if not os.path.exists(dest):
+            os.makedirs(dest)
+
+        # This model is a large model.
+        tf2onnx.convert.from_keras(model, opset=13, output_path=onnx_name)
 
 
 if __name__ == "__main__":

tests/tfhub/tfhub_yamnet_coral.py

@@ -0,0 +1,23 @@
+# 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-coral"
+    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)
+    # WARNING - Error loading model into tflite interpreter: Encountered unresolved custom op: edgetpu-custom-op.Node number 14 (edgetpu-custom-op) failed to prepare.
+    # WARNING - Could not parse attributes for custom op 'TFL_edgetpu-custom-op': 'utf-8' codec can't decode byte 0xc8 in position 0: invalid continuation byte
+    # WARNING - For now, onnxruntime only support float32 type for Gemm rewriter
+    # ERROR - Tensorflow op [tower0/network/layer32/final_output1_prequant: TFL_edgetpu-custom-op] is not supported
+    # ERROR - Unsupported ops: Counter({'TFL_edgetpu-custom-op': 1})
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_yamnet_tf.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/yamnet/1?tf-hub-format=compressed"
+    dest = "tf-yamnet-tf"
+    name = "yamnet"
+    onnx_name = os.path.join(dest, "%s-%d.onnx" % (name, opset))
+
+    imgs = generate_random_images(shape=(16000, ), dtype=numpy.float32, scale=0.)
+
+    benchmark(url, dest, onnx_name, opset, imgs)
+
+
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_yamnet.py renamed to tests/tfhub/tfhub_yamnet_tflite.py

@@ -5,14 +5,16 @@
 
 
 def main(opset=13):
-    url = "https://tfhub.dev/google/coral-model/yamnet/classification/coral/1?coral-format=tflite"
-    dest = "tf-yamnet"
+    url = "https://tfhub.dev/google/lite-model/yamnet/classification/tflite/1?lite-format=tflite"
+    dest = "tf-yamnet-tflite"
     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)
+    # WARNING - For now, onnxruntime only support float32 type for Gemm rewriter
+    # onnxruntime: Could not find an implementation for the node pre_tower/split_prequant:Split(13)
 
 
 if __name__ == "__main__":