| import argparse | |
| import os | |
| import pathlib | |
| import yaml | |
| from pytorch_lightning import Trainer | |
| from pytorch_lightning.callbacks import ModelCheckpoint | |
| from pytorch_lightning.loggers.csv_logs import CSVLogger | |
| from pytorch_lightning.loggers import TensorBoardLogger | |
| from dataset import DataModule | |
| from lightning_module import ( | |
| PretrainLightningModule, | |
| SSLStepLightningModule, | |
| SSLDualLightningModule, | |
| ) | |
| def get_arg(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--config_path", required=True, type=pathlib.Path) | |
| parser.add_argument("--ckpt_path", required=True, type=pathlib.Path) | |
| parser.add_argument( | |
| "--stage", required=True, type=str, choices=["pretrain", "ssl-step", "ssl-dual"] | |
| ) | |
| parser.add_argument("--run_name", required=True, type=str) | |
| return parser.parse_args() | |
| def eval(args, config, output_path): | |
| csvlogger = CSVLogger(save_dir=output_path, name="test_log") | |
| trainer = Trainer( | |
| gpus=-1, | |
| deterministic=False, | |
| auto_select_gpus=True, | |
| benchmark=True, | |
| logger=[csvlogger], | |
| default_root_dir=os.getcwd(), | |
| ) | |
| if config["general"]["stage"] == "pretrain": | |
| model = PretrainLightningModule(config).load_from_checkpoint( | |
| checkpoint_path=args.ckpt_path, config=config | |
| ) | |
| elif config["general"]["stage"] == "ssl-step": | |
| model = SSLStepLightningModule(config).load_from_checkpoint( | |
| checkpoint_path=args.ckpt_path, config=config | |
| ) | |
| elif config["general"]["stage"] == "ssl-dual": | |
| model = SSLDualLightningModule(config).load_from_checkpoint( | |
| checkpoint_path=args.ckpt_path, config=config | |
| ) | |
| else: | |
| raise NotImplementedError() | |
| datamodule = DataModule(config) | |
| trainer.test(model=model, verbose=True, datamodule=datamodule) | |
| if __name__ == "__main__": | |
| args = get_arg() | |
| config = yaml.load(open(args.config_path, "r"), Loader=yaml.FullLoader) | |
| output_path = str(pathlib.Path(config["general"]["output_path"]) / args.run_name) | |
| config["general"]["stage"] = str(getattr(args, "stage")) | |
| eval(args, config, output_path) | |