train.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. # Copyright (c) OpenMMLab. All rights reserved.
  2. import argparse
  3. import logging
  4. import os
  5. import os.path as osp
  6. from mmengine.config import Config, DictAction
  7. from mmengine.logging import print_log
  8. from mmengine.registry import RUNNERS
  9. from mmengine.runner import Runner
  10. from mmdet.utils import setup_cache_size_limit_of_dynamo
  11. def parse_args():
  12. parser = argparse.ArgumentParser(description='Train a detector')
  13. parser.add_argument('config', help='train config file path')
  14. parser.add_argument('--work-dir', help='the dir to save logs and models')
  15. parser.add_argument(
  16. '--amp',
  17. action='store_true',
  18. default=False,
  19. help='enable automatic-mixed-precision training')
  20. parser.add_argument(
  21. '--auto-scale-lr',
  22. action='store_true',
  23. help='enable automatically scaling LR.')
  24. parser.add_argument(
  25. '--resume',
  26. nargs='?',
  27. type=str,
  28. const='auto',
  29. help='If specify checkpoint path, resume from it, while if not '
  30. 'specify, try to auto resume from the latest checkpoint '
  31. 'in the work directory.')
  32. parser.add_argument(
  33. '--cfg-options',
  34. nargs='+',
  35. action=DictAction,
  36. help='override some settings in the used config, the key-value pair '
  37. 'in xxx=yyy format will be merged into config file. If the value to '
  38. 'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
  39. 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
  40. 'Note that the quotation marks are necessary and that no white space '
  41. 'is allowed.')
  42. parser.add_argument(
  43. '--launcher',
  44. choices=['none', 'pytorch', 'slurm', 'mpi'],
  45. default='none',
  46. help='job launcher')
  47. # When using PyTorch version >= 2.0.0, the `torch.distributed.launch`
  48. # will pass the `--local-rank` parameter to `tools/train.py` instead
  49. # of `--local_rank`.
  50. parser.add_argument('--local_rank', '--local-rank', type=int, default=0)
  51. args = parser.parse_args()
  52. if 'LOCAL_RANK' not in os.environ:
  53. os.environ['LOCAL_RANK'] = str(args.local_rank)
  54. return args
  55. def main():
  56. args = parse_args()
  57. # Reduce the number of repeated compilations and improve
  58. # training speed.
  59. setup_cache_size_limit_of_dynamo()
  60. # load config
  61. cfg = Config.fromfile(args.config)
  62. cfg.launcher = args.launcher
  63. if args.cfg_options is not None:
  64. cfg.merge_from_dict(args.cfg_options)
  65. # work_dir is determined in this priority: CLI > segment in file > filename
  66. if args.work_dir is not None:
  67. # update configs according to CLI args if args.work_dir is not None
  68. cfg.work_dir = args.work_dir
  69. elif cfg.get('work_dir', None) is None:
  70. # use config filename as default work_dir if cfg.work_dir is None
  71. cfg.work_dir = osp.join('./work_dirs',
  72. osp.splitext(osp.basename(args.config))[0])
  73. # enable automatic-mixed-precision training
  74. if args.amp is True:
  75. optim_wrapper = cfg.optim_wrapper.type
  76. if optim_wrapper == 'AmpOptimWrapper':
  77. print_log(
  78. 'AMP training is already enabled in your config.',
  79. logger='current',
  80. level=logging.WARNING)
  81. else:
  82. assert optim_wrapper == 'OptimWrapper', (
  83. '`--amp` is only supported when the optimizer wrapper type is '
  84. f'`OptimWrapper` but got {optim_wrapper}.')
  85. cfg.optim_wrapper.type = 'AmpOptimWrapper'
  86. cfg.optim_wrapper.loss_scale = 'dynamic'
  87. # enable automatically scaling LR
  88. if args.auto_scale_lr:
  89. if 'auto_scale_lr' in cfg and \
  90. 'enable' in cfg.auto_scale_lr and \
  91. 'base_batch_size' in cfg.auto_scale_lr:
  92. cfg.auto_scale_lr.enable = True
  93. else:
  94. raise RuntimeError('Can not find "auto_scale_lr" or '
  95. '"auto_scale_lr.enable" or '
  96. '"auto_scale_lr.base_batch_size" in your'
  97. ' configuration file.')
  98. # resume is determined in this priority: resume from > auto_resume
  99. if args.resume == 'auto':
  100. cfg.resume = True
  101. cfg.load_from = None
  102. elif args.resume is not None:
  103. cfg.resume = True
  104. cfg.load_from = args.resume
  105. # build the runner from config
  106. if 'runner_type' not in cfg:
  107. # build the default runner
  108. runner = Runner.from_cfg(cfg)
  109. else:
  110. # build customized runner from the registry
  111. # if 'runner_type' is set in the cfg
  112. runner = RUNNERS.build(cfg)
  113. # start training
  114. runner.train()
  115. if __name__ == '__main__':
  116. main()