File size: 5,942 Bytes
4947683 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | from pathlib import Path
from typing import List
import sys
sys.path.append('.')
import hydra
import numpy as np
import torch
import omegaconf
import pytorch_lightning as pl
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig, OmegaConf
from pytorch_lightning import seed_everything, Callback
from pytorch_lightning.callbacks import (
EarlyStopping,
LearningRateMonitor,
ModelCheckpoint,
)
from pytorch_lightning.loggers import WandbLogger
from diffcsp.common.utils import log_hyperparameters, PROJECT_ROOT
import wandb
def build_callbacks(cfg: DictConfig) -> List[Callback]:
callbacks: List[Callback] = []
if "lr_monitor" in cfg.logging:
hydra.utils.log.info("Adding callback <LearningRateMonitor>")
callbacks.append(
LearningRateMonitor(
logging_interval=cfg.logging.lr_monitor.logging_interval,
log_momentum=cfg.logging.lr_monitor.log_momentum,
)
)
if "early_stopping" in cfg.train:
hydra.utils.log.info("Adding callback <EarlyStopping>")
callbacks.append(
EarlyStopping(
monitor=cfg.train.monitor_metric,
mode=cfg.train.monitor_metric_mode,
patience=cfg.train.early_stopping.patience,
verbose=cfg.train.early_stopping.verbose,
)
)
if "model_checkpoints" in cfg.train:
hydra.utils.log.info("Adding callback <ModelCheckpoint>")
callbacks.append(
ModelCheckpoint(
dirpath=Path(HydraConfig.get().run.dir),
monitor=cfg.train.monitor_metric,
mode=cfg.train.monitor_metric_mode,
save_top_k=cfg.train.model_checkpoints.save_top_k,
verbose=cfg.train.model_checkpoints.verbose,
save_last=cfg.train.model_checkpoints.save_last,
)
)
return callbacks
def run(cfg: DictConfig) -> None:
"""
Generic train loop
:param cfg: run configuration, defined by Hydra in /conf
"""
if cfg.train.deterministic:
seed_everything(cfg.train.random_seed)
if cfg.train.pl_trainer.fast_dev_run:
hydra.utils.log.info(
f"Debug mode <{cfg.train.pl_trainer.fast_dev_run=}>. "
f"Forcing debugger friendly configuration!"
)
# Debuggers don't like GPUs nor multiprocessing
cfg.train.pl_trainer.gpus = 0
cfg.data.datamodule.num_workers.train = 0
cfg.data.datamodule.num_workers.val = 0
cfg.data.datamodule.num_workers.test = 0
# Switch wandb mode to offline to prevent online logging
cfg.logging.wandb.mode = "offline"
# Hydra run directory
hydra_dir = Path(HydraConfig.get().run.dir)
# Instantiate datamodule
hydra.utils.log.info(f"Instantiating <{cfg.data.datamodule._target_}>")
datamodule: pl.LightningDataModule = hydra.utils.instantiate(
cfg.data.datamodule, _recursive_=False
)
# Instantiate model
hydra.utils.log.info(f"Instantiating <{cfg.model._target_}>")
model: pl.LightningModule = hydra.utils.instantiate(
cfg.model,
optim=cfg.optim,
data=cfg.data,
logging=cfg.logging,
_recursive_=False,
)
# Pass scaler from datamodule to model
hydra.utils.log.info(f"Passing scaler from datamodule to model <{datamodule.scaler}>")
if datamodule.scaler is not None:
model.lattice_scaler = datamodule.lattice_scaler.copy()
model.scaler = datamodule.scaler.copy()
torch.save(datamodule.lattice_scaler, hydra_dir / 'lattice_scaler.pt')
torch.save(datamodule.scaler, hydra_dir / 'prop_scaler.pt')
# Instantiate the callbacks
callbacks: List[Callback] = build_callbacks(cfg=cfg)
# Logger instantiation/configuration
wandb_logger = None
if "wandb" in cfg.logging:
hydra.utils.log.info("Instantiating <WandbLogger>")
wandb_config = cfg.logging.wandb
wandb_logger = WandbLogger(
**wandb_config,
settings=wandb.Settings(start_method="fork"),
tags=cfg.core.tags,
)
hydra.utils.log.info("W&B is now watching <{cfg.logging.wandb_watch.log}>!")
wandb_logger.watch(
model,
log=cfg.logging.wandb_watch.log,
log_freq=cfg.logging.wandb_watch.log_freq,
)
# Store the YaML config separately into the wandb dir
yaml_conf: str = OmegaConf.to_yaml(cfg=cfg)
(hydra_dir / "hparams.yaml").write_text(yaml_conf)
# Load checkpoint (if exist)
ckpts = list(hydra_dir.glob('*.ckpt'))
if len(ckpts) > 0:
ckpt_epochs = np.array([int(ckpt.parts[-1].split('-')[0].split('=')[1]) for ckpt in ckpts])
ckpt = str(ckpts[ckpt_epochs.argsort()[-1]])
hydra.utils.log.info(f"found checkpoint: {ckpt}")
else:
ckpt = None
hydra.utils.log.info("Instantiating the Trainer")
trainer = pl.Trainer(
default_root_dir=hydra_dir,
logger=wandb_logger,
callbacks=callbacks,
deterministic=cfg.train.deterministic,
check_val_every_n_epoch=cfg.logging.val_check_interval,
progress_bar_refresh_rate=cfg.logging.progress_bar_refresh_rate,
resume_from_checkpoint=ckpt,
**cfg.train.pl_trainer,
)
log_hyperparameters(trainer=trainer, model=model, cfg=cfg)
hydra.utils.log.info("Starting training!")
trainer.fit(model=model, datamodule=datamodule)
hydra.utils.log.info("Starting testing!")
trainer.test(datamodule=datamodule)
# Logger closing to release resources/avoid multi-run conflicts
if wandb_logger is not None:
wandb_logger.experiment.finish()
@hydra.main(config_path=str(PROJECT_ROOT / "conf"), config_name="default", version_base="1.1")
def main(cfg: omegaconf.DictConfig):
run(cfg)
if __name__ == "__main__":
main()
|