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