Spaces:
Running
Running
| import os | |
| import os.path as osp | |
| import time | |
| import sys | |
| CODE_SPACE=os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| sys.path.append(CODE_SPACE) | |
| #os.chdir(CODE_SPACE) | |
| import argparse | |
| import mmcv | |
| import torch | |
| import torch.distributed as dist | |
| import torch.multiprocessing as mp | |
| try: | |
| from mmcv.utils import Config, DictAction | |
| except: | |
| from mmengine import Config, DictAction | |
| from datetime import timedelta | |
| import random | |
| import numpy as np | |
| from mono.datasets.distributed_sampler import log_canonical_transfer_info | |
| from mono.utils.comm import init_env | |
| from mono.utils.logger import setup_logger | |
| from mono.utils.db import load_data_info, reset_ckpt_path | |
| from mono.model.monodepth_model import get_configured_monodepth_model | |
| from mono.datasets.distributed_sampler import build_dataset_n_sampler_with_cfg | |
| from mono.utils.running import load_ckpt | |
| from mono.utils.do_test import do_test_with_dataloader, do_test_check_data | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description='Train a segmentor') | |
| parser.add_argument('config', help='train config file path') | |
| parser.add_argument('--show-dir', help='the dir to save logs and visualization results') | |
| parser.add_argument( | |
| '--load-from', help='the checkpoint file to load weights from') | |
| parser.add_argument('--node_rank', type=int, default=0) | |
| parser.add_argument('--nnodes', | |
| type=int, | |
| default=1, | |
| help='number of nodes') | |
| parser.add_argument( | |
| '--options', nargs='+', action=DictAction, help='custom options') | |
| parser.add_argument( | |
| '--launcher', choices=['None', 'pytorch', 'slurm'], default='slurm', | |
| help='job launcher') | |
| args = parser.parse_args() | |
| return args | |
| def main(args): | |
| os.chdir(CODE_SPACE) | |
| cfg = Config.fromfile(args.config) | |
| cfg.dist_params.nnodes = args.nnodes | |
| cfg.dist_params.node_rank = args.node_rank | |
| if args.options is not None: | |
| cfg.merge_from_dict(args.options) | |
| # set cudnn_benchmark | |
| #if cfg.get('cudnn_benchmark', False) and args.launcher != 'ror': | |
| # torch.backends.cudnn.benchmark = True | |
| # show_dir is determined in this priority: CLI > segment in file > filename | |
| if args.show_dir is not None: | |
| # update configs according to CLI args if args.show_dir is not None | |
| cfg.show_dir = args.show_dir | |
| elif cfg.get('show_dir', None) is None: | |
| # use config filename + timestamp as default show_dir if cfg.show_dir is None | |
| cfg.show_dir = osp.join('./show_dirs', | |
| osp.splitext(osp.basename(args.config))[0], | |
| args.timestamp) | |
| # ckpt path | |
| if args.load_from is None: | |
| raise RuntimeError('Please set model path!') | |
| cfg.load_from = args.load_from | |
| # create show dir | |
| os.makedirs(osp.abspath(cfg.show_dir), exist_ok=True) | |
| # init the logger before other steps | |
| cfg.log_file = osp.join(cfg.show_dir, f'{args.timestamp}.log') | |
| logger = setup_logger(cfg.log_file) | |
| # log some basic info | |
| logger.info(f'Config:\n{cfg.pretty_text}') | |
| # load db_info for data | |
| # load data info | |
| data_info = {} | |
| load_data_info('data_server_info', data_info=data_info) | |
| cfg.db_info = data_info | |
| # update check point info | |
| reset_ckpt_path(cfg.model, data_info) | |
| # log data transfer to canonical space info | |
| # log_canonical_transfer_info(cfg) | |
| # init distributed env first, since logger depends on the dist info. | |
| if args.launcher == 'none': | |
| cfg.distributed = False | |
| else: | |
| cfg.distributed = True | |
| init_env(args.launcher, cfg) | |
| logger.info(f'Distributed training: {cfg.distributed}') | |
| # dump config | |
| cfg.dump(osp.join(cfg.show_dir, osp.basename(args.config))) | |
| if not cfg.distributed: | |
| main_worker(0, cfg, args.launcher) | |
| else: | |
| mp.spawn(main_worker, nprocs=cfg.dist_params.num_gpus_per_node, args=(cfg, args.launcher)) | |
| def main_worker(local_rank: int, cfg: dict, launcher: str): | |
| if cfg.distributed: | |
| cfg.dist_params.global_rank = cfg.dist_params.node_rank * cfg.dist_params.num_gpus_per_node + local_rank | |
| cfg.dist_params.local_rank = local_rank | |
| torch.cuda.set_device(local_rank) | |
| default_timeout = timedelta(minutes=30) | |
| dist.init_process_group(backend=cfg.dist_params.backend, | |
| init_method=cfg.dist_params.dist_url, | |
| world_size=cfg.dist_params.world_size, | |
| rank=cfg.dist_params.global_rank, | |
| timeout=default_timeout,) | |
| logger = setup_logger(cfg.log_file) | |
| # build model | |
| model = get_configured_monodepth_model(cfg, | |
| None, | |
| ) | |
| # build datasets | |
| test_dataset, test_sampler = build_dataset_n_sampler_with_cfg(cfg, 'test') | |
| # build data loaders | |
| test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, | |
| batch_size=1, | |
| num_workers=1, | |
| sampler=test_sampler, | |
| drop_last=False) | |
| # config distributed training | |
| if cfg.distributed: | |
| model = torch.nn.parallel.DistributedDataParallel(model.cuda(), | |
| device_ids=[local_rank], | |
| output_device=local_rank, | |
| find_unused_parameters=True) | |
| else: | |
| model = torch.nn.DataParallel(model.cuda()) | |
| # load ckpt | |
| #model, _, _, _ = load_ckpt(cfg.load_from, model, strict_match=False) | |
| model.eval() | |
| do_test_with_dataloader(model, cfg, test_dataloader, logger=logger, is_distributed=cfg.distributed) | |
| # do_test_check_data(model, cfg, test_dataloader, logger=logger, is_distributed=cfg.distributed, local_rank=local_rank) | |
| if __name__=='__main__': | |
| # load args | |
| args = parse_args() | |
| timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) | |
| args.timestamp = timestamp | |
| main(args) | |