train.py

import hydra
from hydra.utils import instantiate
from omegaconf import DictConfig, OmegaConf
import pytorch_lightning as pl


@hydra.main(config_path='conf', config_name='config')
def main(cfg: DictConfig) -> None:

    pl._logger.handlers = []
    pl._logger.propagate = True

    if cfg.seed is not None:
        pl.seed_everything(cfg.seed)

    model = instantiate(cfg.pipeline, cfg=cfg)

    cfg_trainer = dict(cfg.pl_trainer)
    if cfg.logging:
        loggers = []
        for cfg_log in cfg.logging:
            loggers.append(instantiate(cfg_log))
        cfg_trainer['logger'] = loggers
    if cfg.callbacks:
        callbacks = []
        for cfg_callback in cfg.callbacks:
            callbacks.append(instantiate(cfg_callback))
        cfg_trainer['callbacks'] =callbacks

    trainer = pl.Trainer(**cfg_trainer)
    datamodule = instantiate(cfg.dataset)
    trainer.fit(model, datamodule)

    if cfg.run_test:
        trainer.test(model, datamodule=datamodule)


if __name__ == '__main__':
    main()