onnx_config.py 641 B

123456789101112131415161718192021222324
  1. import os
  2. def get_op_num_threads(env_name: str) -> int:
  3. env_value = os.getenv(env_name, None)
  4. return get_op_num_threads_from_value(env_value)
  5. def get_op_num_threads_from_value(env_value: str) -> int:
  6. if env_value is not None:
  7. try:
  8. num_threads = int(env_value)
  9. if num_threads > 0:
  10. return num_threads
  11. except ValueError:
  12. return -1
  13. return -1
  14. if __name__ == '__main__':
  15. print(get_op_num_threads_from_value('1'))
  16. print(get_op_num_threads_from_value('0'))
  17. print(get_op_num_threads_from_value('-1'))
  18. print(get_op_num_threads_from_value('abc'))