File size: 2,013 Bytes
14ae0ea
a89496d
 
ccecb22
8949a8c
a89496d
14ae0ea
 
7bb4fe3
a89496d
 
1530829
 
a89496d
 
 
 
8949a8c
e8eaf47
 
ace4057
 
 
 
 
 
 
 
 
 
e8eaf47
a89496d
 
 
 
 
 
 
8949a8c
a89496d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5570d2c
14ae0ea
 
 
 
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
import pytorch_lightning as pl
import hydra
from omegaconf import DictConfig
import remfx.utils as utils

log = utils.get_logger(__name__)


@hydra.main(version_base=None, config_path="../cfg", config_name="config.yaml")
def main(cfg: DictConfig):
    # Apply seed for reproducibility
    if cfg.seed:
        pl.seed_everything(cfg.seed)
    log.info(f"Instantiating datamodule <{cfg.datamodule._target_}>.")
    datamodule = hydra.utils.instantiate(cfg.datamodule, _convert_="partial")
    log.info(f"Instantiating model <{cfg.model._target_}>.")
    model = hydra.utils.instantiate(cfg.model, _convert_="partial")

    if "ckpt_path" in cfg:
        log.info(f"Loading checkpoint from <{cfg.ckpt_path}>.")
        model.load_from_checkpoint(
            cfg.ckpt_path,
            lr=model.lr,
            lr_beta1=model.lr_beta1,
            lr_beta2=model.lr_beta2,
            lr_eps=model.lr_eps,
            lr_weight_decay=model.lr_weight_decay,
            sample_rate=model.sample_rate,
            network=model.model,
        )

    # Init all callbacks
    callbacks = []
    if "callbacks" in cfg:
        for _, cb_conf in cfg["callbacks"].items():
            if "_target_" in cb_conf:
                log.info(f"Instantiating callback <{cb_conf._target_}>.")
                callbacks.append(hydra.utils.instantiate(cb_conf, _convert_="partial"))

    logger = hydra.utils.instantiate(cfg.logger, _convert_="partial")
    log.info(f"Instantiating trainer <{cfg.trainer._target_}>.")
    trainer = hydra.utils.instantiate(
        cfg.trainer, callbacks=callbacks, logger=logger, _convert_="partial"
    )
    log.info("Logging hyperparameters!")
    utils.log_hyperparameters(
        config=cfg,
        model=model,
        datamodule=datamodule,
        trainer=trainer,
        callbacks=callbacks,
        logger=logger,
    )
    trainer.fit(model=model, datamodule=datamodule)
    trainer.test(model=model, datamodule=datamodule, ckpt_path="best")


if __name__ == "__main__":
    main()