config.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. # time: 2023/3/2 10:28
  4. # file: config.py
  5. # author: David
  6. # company: shenyang JY
  7. """
  8. 模型调参及系统功能配置
  9. """
  10. import os, threading
  11. import argparse, yaml
  12. class myargparse(argparse.ArgumentParser):
  13. _save_lock = threading.Lock()
  14. def __init__(self, description, add_help):
  15. super(myargparse, self).__init__(description=description, add_help=add_help)
  16. self.add_argument(
  17. '-c',
  18. '--config_yaml',
  19. default=
  20. 'config.yml',
  21. type=str,
  22. metavar='FILE',
  23. help='YAML config file specifying default arguments')
  24. def _parse_args_and_yaml(self):
  25. given_configs, remaining = self.parse_known_args()
  26. current_path = os.path.dirname(__file__)
  27. if given_configs.config_yaml:
  28. with open(current_path + '/' + given_configs.config_yaml, 'r', encoding='utf-8') as f:
  29. cfg = yaml.safe_load(f)
  30. self.set_defaults(**cfg)
  31. # defaults will have been overridden if config file specified.
  32. opt = self.parse_args(remaining)
  33. # Cache the args as a text string to save them in the output dir later
  34. opt_text = yaml.safe_dump(opt.__dict__, default_flow_style=False)
  35. return opt, opt_text
  36. def parse_args_and_yaml(self):
  37. return self._parse_args_and_yaml()[0]
  38. def save_args_yml(self, opt):
  39. current_path = os.path.dirname(__file__)
  40. with myargparse._save_lock:
  41. file_path = os.path.join(current_path, 'config.yml')
  42. with open(file_path, mode='w', encoding='utf-8') as f:
  43. yaml.safe_dump(vars(opt), f)
  44. if __name__ == '__main__':
  45. args = myargparse(discription="场站端配置", add_help=False)
  46. opt = args.parse_args_and_yaml()