procgne/Plonk

1import os2from models.module import DiffGeolocalizer3import hydra4import wandb5from os.path import isfile, join6from shutil import copyfile7 8import torch9 10from omegaconf import OmegaConf11from omegaconf import open_dict12from hydra.core.hydra_config import HydraConfig13from hydra.utils import instantiate14from pytorch_lightning.callbacks import LearningRateMonitor15from lightning_fabric.utilities.rank_zero import _get_rank16 17from models.module import DiffGeolocalizer18 19torch.set_float32_matmul_precision("high")  # TODO do we need that?20 21# Registering the "eval" resolver allows for advanced config22# interpolation with arithmetic operations in hydra:23# https://omegaconf.readthedocs.io/en/2.3_branch/how_to_guides.html24OmegaConf.register_new_resolver("eval", eval)25 26 27def load_model(cfg, dict_config, wandb_id):28    logger = instantiate(cfg.logger, id=open(wandb_id, "r").read(), resume="allow")29    model = DiffGeolocalizer.load_from_checkpoint(cfg.checkpoint, cfg=cfg.model)30    trainer = instantiate(cfg.trainer, strategy=cfg.trainer.strategy, logger=logger)31    return trainer, model32 33 34def hydra_boilerplate(cfg):35    dict_config = OmegaConf.to_container(cfg, resolve=True)36    trainer, model = load_model(cfg, dict_config, cfg.wandb_id)37    return trainer, model38 39 40import copy41 42 43def generate_datamodules(cfg_):44    for f in os.listdir(cfg_.test_dir):45        cfg = copy.deepcopy(cfg_)46        # open join(f, directory) with OmegaConf47        with open_dict(cfg):48            cfg_new = OmegaConf.load(join(cfg.test_dir, f))49            cfg.datamodule = cfg_new.datamodule50            cfg.dataset = cfg_new.dataset51            cfg.dataset.test_transform = cfg_.dataset.test_transform52 53        datamodule = instantiate(cfg.datamodule)54        yield datamodule55 56 57if __name__ == "__main__":58    import sys59 60    sys.argv = (61        [sys.argv[0]]62        + ["+pt_model_path=${hydra:runtime.config_sources}"]63        + sys.argv[1:]64    )65 66    @hydra.main(version_base=None)67    def main(cfg):68        # print(hydra.runtime.config_sources)69        with open_dict(cfg):70            path = cfg.pt_model_path[1]["path"]71            cfg.wandb_id = join(path, "wandb_id.txt")72            cfg.checkpoint = join(path, "last.ckpt")73            cfg.computer.devices = 174 75        (76            trainer,77            model,78        ) = hydra_boilerplate(cfg)79        for datamodule in generate_datamodules(cfg):80            model.datamodule = datamodule81            model.datamodule.setup()82            print("Testing on", datamodule.test_dataset.class_name)83            trainer.test(model, datamodule=datamodule)84 85    main()86