Spaces:
Running
Running
| import pytorch_lightning as pl | |
| import argparse | |
| import pprint | |
| from loguru import logger as loguru_logger | |
| from src.config.default import get_cfg_defaults | |
| from src.utils.profiler import build_profiler | |
| from src.lightning_trainer.data import MultiSceneDataModule | |
| from src.lightning_trainer.trainer import PL_Trainer | |
| def parse_args(): | |
| # init a costum parser which will be added into pl.Trainer parser | |
| # check documentation: https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html#trainer-flags | |
| parser = argparse.ArgumentParser( | |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter | |
| ) | |
| parser.add_argument("data_cfg_path", type=str, help="data config path") | |
| parser.add_argument("main_cfg_path", type=str, help="main config path") | |
| parser.add_argument( | |
| "--ckpt_path", | |
| type=str, | |
| default="weights/indoor_ds.ckpt", | |
| help="path to the checkpoint", | |
| ) | |
| parser.add_argument( | |
| "--dump_dir", | |
| type=str, | |
| default=None, | |
| help="if set, the matching results will be dump to dump_dir", | |
| ) | |
| parser.add_argument( | |
| "--profiler_name", | |
| type=str, | |
| default=None, | |
| help="options: [inference, pytorch], or leave it unset", | |
| ) | |
| parser.add_argument("--batch_size", type=int, default=1, help="batch_size per gpu") | |
| parser.add_argument("--num_workers", type=int, default=2) | |
| parser.add_argument( | |
| "--thr", | |
| type=float, | |
| default=None, | |
| help="modify the coarse-level matching threshold.", | |
| ) | |
| parser = pl.Trainer.add_argparse_args(parser) | |
| return parser.parse_args() | |
| if __name__ == "__main__": | |
| # parse arguments | |
| args = parse_args() | |
| pprint.pprint(vars(args)) | |
| # init default-cfg and merge it with the main- and data-cfg | |
| config = get_cfg_defaults() | |
| config.merge_from_file(args.main_cfg_path) | |
| config.merge_from_file(args.data_cfg_path) | |
| pl.seed_everything(config.TRAINER.SEED) # reproducibility | |
| # tune when testing | |
| if args.thr is not None: | |
| config.MODEL.MATCH_COARSE.THR = args.thr | |
| loguru_logger.info(f"Args and config initialized!") | |
| # lightning module | |
| profiler = build_profiler(args.profiler_name) | |
| model = PL_Trainer( | |
| config, | |
| pretrained_ckpt=args.ckpt_path, | |
| profiler=profiler, | |
| dump_dir=args.dump_dir, | |
| ) | |
| loguru_logger.info(f"Model-lightning initialized!") | |
| # lightning data | |
| data_module = MultiSceneDataModule(args, config) | |
| loguru_logger.info(f"DataModule initialized!") | |
| # lightning trainer | |
| trainer = pl.Trainer.from_argparse_args( | |
| args, replace_sampler_ddp=False, logger=False | |
| ) | |
| loguru_logger.info(f"Start testing!") | |
| trainer.test(model, datamodule=data_module, verbose=False) | |