| 92 | |
| 93 | |
| 94 | def main(): |
| 95 | args = parse_args() |
| 96 | |
| 97 | cfg = Config.fromfile(args.config) |
| 98 | if args.cfg_options is not None: |
| 99 | cfg.merge_from_dict(args.cfg_options) |
| 100 | # import modules from string list. |
| 101 | if cfg.get('custom_imports', None): |
| 102 | from mmcv.utils import import_modules_from_strings |
| 103 | import_modules_from_strings(**cfg['custom_imports']) |
| 104 | # set cudnn_benchmark |
| 105 | if cfg.get('cudnn_benchmark', False): |
| 106 | torch.backends.cudnn.benchmark = True |
| 107 | |
| 108 | # work_dir is determined in this priority: CLI > segment in file > filename |
| 109 | if args.work_dir is not None: |
| 110 | # update configs according to CLI args if args.work_dir is not None |
| 111 | cfg.work_dir = args.work_dir |
| 112 | elif cfg.get('work_dir', None) is None: |
| 113 | # use config filename as default work_dir if cfg.work_dir is None |
| 114 | cfg.work_dir = osp.join('./work_dirs', |
| 115 | osp.splitext(osp.basename(args.config))[0]) |
| 116 | if args.resume_from is not None: |
| 117 | cfg.resume_from = args.resume_from |
| 118 | if args.gpu_ids is not None: |
| 119 | cfg.gpu_ids = args.gpu_ids |
| 120 | else: |
| 121 | cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus) |
| 122 | |
| 123 | # init distributed env first, since logger depends on the dist info. |
| 124 | if args.launcher == 'none': |
| 125 | distributed = False |
| 126 | else: |
| 127 | distributed = True |
| 128 | init_dist(args.launcher, **cfg.dist_params) |
| 129 | # re-set gpu_ids with distributed training mode |
| 130 | _, world_size = get_dist_info() |
| 131 | cfg.gpu_ids = range(world_size) |
| 132 | |
| 133 | # create work_dir |
| 134 | mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) |
| 135 | # dump config |
| 136 | cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config))) |
| 137 | # init the logger before other steps |
| 138 | timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) |
| 139 | log_file = osp.join(cfg.work_dir, f'{timestamp}.log') |
| 140 | logger = get_root_logger(log_file=log_file, log_level=cfg.log_level) |
| 141 | |
| 142 | # init the meta dict to record some important information such as |
| 143 | # environment info and seed, which will be logged |
| 144 | meta = dict() |
| 145 | # log env info |
| 146 | env_info_dict = collect_env() |
| 147 | env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()]) |
| 148 | dash_line = '-' * 60 + '\n' |
| 149 | logger.info('Environment info:\n' + dash_line + env_info + '\n' + |
| 150 | dash_line) |
| 151 | meta['env_info'] = env_info |