| 61 | |
| 62 | |
| 63 | def main(): |
| 64 | args = parse_args() |
| 65 | |
| 66 | # load config |
| 67 | cfg = Config.fromfile(args.config) |
| 68 | |
| 69 | # TODO: We will unify the ceph support approach with other OpenMMLab repos |
| 70 | if args.ceph: |
| 71 | cfg = replace_ceph_backend(cfg) |
| 72 | |
| 73 | cfg.launcher = args.launcher |
| 74 | if args.cfg_options is not None: |
| 75 | cfg.merge_from_dict(args.cfg_options) |
| 76 | |
| 77 | # work_dir is determined in this priority: CLI > segment in file > filename |
| 78 | if args.work_dir is not None: |
| 79 | # update configs according to CLI args if args.work_dir is not None |
| 80 | cfg.work_dir = args.work_dir |
| 81 | elif cfg.get('work_dir', None) is None: |
| 82 | # use config filename as default work_dir if cfg.work_dir is None |
| 83 | cfg.work_dir = osp.join('./work_dirs', |
| 84 | osp.splitext(osp.basename(args.config))[0]) |
| 85 | |
| 86 | # enable automatic-mixed-precision training |
| 87 | if args.amp is True: |
| 88 | optim_wrapper = cfg.optim_wrapper.type |
| 89 | if optim_wrapper == 'AmpOptimWrapper': |
| 90 | print_log( |
| 91 | 'AMP training is already enabled in your config.', |
| 92 | logger='current', |
| 93 | level=logging.WARNING) |
| 94 | else: |
| 95 | assert optim_wrapper == 'OptimWrapper', ( |
| 96 | '`--amp` is only supported when the optimizer wrapper type is ' |
| 97 | f'`OptimWrapper` but got {optim_wrapper}.') |
| 98 | cfg.optim_wrapper.type = 'AmpOptimWrapper' |
| 99 | cfg.optim_wrapper.loss_scale = 'dynamic' |
| 100 | |
| 101 | # enable automatically scaling LR |
| 102 | if args.auto_scale_lr: |
| 103 | if 'auto_scale_lr' in cfg and \ |
| 104 | 'enable' in cfg.auto_scale_lr and \ |
| 105 | 'base_batch_size' in cfg.auto_scale_lr: |
| 106 | cfg.auto_scale_lr.enable = True |
| 107 | else: |
| 108 | raise RuntimeError('Can not find "auto_scale_lr" or ' |
| 109 | '"auto_scale_lr.enable" or ' |
| 110 | '"auto_scale_lr.base_batch_size" in your' |
| 111 | ' configuration file.') |
| 112 | |
| 113 | # resume is determined in this priority: resume from > auto_resume |
| 114 | if args.resume == 'auto': |
| 115 | cfg.resume = True |
| 116 | cfg.load_from = None |
| 117 | elif args.checkpoint is not None: |
| 118 | cfg.load_from = args.checkpoint |
| 119 | elif args.resume is not None: |
| 120 | cfg.resume = True |