procgne/Plonk

1import os2from models.module import DiffGeolocalizer3import hydra4from os.path import join5 6import torch7 8from omegaconf import OmegaConf9from omegaconf import open_dict10from hydra.utils import instantiate11 12from models.eval_best_model import EvalModule13 14torch.set_float32_matmul_precision("high")15 16# Registering the "eval" resolver allows for advanced config17# interpolation with arithmetic operations in hydra:18# https://omegaconf.readthedocs.io/en/2.3_branch/how_to_guides.html19OmegaConf.register_new_resolver("eval", eval)20 21 22def load_model(cfg, dict_config, wandb_id):23    logger = instantiate(cfg.logger, id=wandb_id, resume="allow")24    log_dict = {"model": dict_config["model"], "dataset": dict_config["dataset"]}25    logger._wandb_init.update({"config": log_dict})26    model = EvalModule(cfg.model)27    trainer = instantiate(28        cfg.trainer, strategy=cfg.trainer.strategy29    )  # , logger=logger)30    return trainer, model31 32 33def hydra_boilerplate(cfg):34    dict_config = OmegaConf.to_container(cfg, resolve=True)35    trainer, model = load_model(cfg, dict_config, cfg.wandb_id)36    return trainer, model37 38 39import copy40 41 42def init_datamodule(cfg):43    datamodule = instantiate(cfg.datamodule)44    return datamodule45 46 47if __name__ == "__main__":48    import sys49 50    sys.argv = (51        [sys.argv[0]]52        + ["+pt_model_path=${hydra:runtime.config_sources}"]53        + sys.argv[1:]54    )55 56    @hydra.main(config_path="configs", config_name="config", version_base=None)57    def main(cfg):58        # print(hydra.runtime.config_sources)59        with open_dict(cfg):60            path = cfg.pt_model_path[1]["path"]61            cfg.wandb_id = join(path, "wandb_id.txt")62            cfg.checkpoint = join(path, "last.ckpt")63            cfg.computer.devices = 164 65        (66            trainer,67            model,68        ) = hydra_boilerplate(cfg)69        datamodule = init_datamodule(cfg)70        trainer.test(model, datamodule=datamodule)71 72    main()73