Spaces:
Running
Running
File size: 2,791 Bytes
404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
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)
|