refactor

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

Author: sdpython

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

@@ -15,4 +15,5 @@ def main(opset=13):
     benchmark(url, dest, onnx_name, opset, imgs)
 
 
-main()
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_resnet_v2_152.py

@@ -15,4 +15,5 @@ def main(opset=13):
     benchmark(url, dest, onnx_name, opset, imgs)
 
 
-main()
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_spam_detection.py

@@ -16,4 +16,5 @@ def main(opset=13):
     benchmark(url, dest, onnx_name, opset, imgs)
 
 
-main()
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_thunder.py

@@ -16,4 +16,5 @@ def main(opset=13):
               signature='serving_default')
 
 
-main()
+if __name__ == "__main__":
+    main()

tests/tfhub/tfhub_yamnet.py

@@ -15,4 +15,5 @@ def main(opset=13):
     benchmark_tflite(url, dest, onnx_name, opset, imgs)
 
 
-main()
+if __name__ == "__main__":
+    main()