check_onnx_model_mobile_usability.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. # Copyright (c) Microsoft Corporation. All rights reserved.
  2. # Licensed under the MIT License.
  3. import argparse
  4. import logging
  5. import pathlib
  6. # need this before the mobile helper imports for some reason
  7. logging.basicConfig(format="%(levelname)s: %(message)s")
  8. from .mobile_helpers import usability_checker # noqa: E402
  9. def check_usability():
  10. parser = argparse.ArgumentParser(
  11. description="""Analyze an ONNX model to determine how well it will work in mobile scenarios.""",
  12. formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  13. )
  14. parser.add_argument("--log_level", choices=["debug", "info"], default="info", help="Logging level")
  15. parser.add_argument("model_path", help="Path to ONNX model to check", type=pathlib.Path)
  16. args = parser.parse_args()
  17. logger = logging.getLogger("check_usability")
  18. if args.log_level == "debug":
  19. logger.setLevel(logging.DEBUG)
  20. elif args.log_level == "info":
  21. logger.setLevel(logging.INFO)
  22. elif args.log_level == "warning":
  23. logger.setLevel(logging.WARNING)
  24. else:
  25. logger.setLevel(logging.ERROR)
  26. try_eps = usability_checker.analyze_model(args.model_path, skip_optimize=False, logger=logger)
  27. if try_eps:
  28. logger.info(
  29. "As NNAPI or CoreML may provide benefits with this model it is recommended to compare the "
  30. "performance of the model using the NNAPI EP on Android, and the CoreML EP on iOS, "
  31. "against the performance using the CPU EP."
  32. )
  33. else:
  34. logger.info("For optimal performance the model should be used with the CPU EP. ")
  35. if __name__ == "__main__":
  36. check_usability()