| import os |
|
|
| import hydra |
| import omegaconf |
| import pytorch_lightning as pl |
| import torch |
| from pytorch_lightning.loggers import TensorBoardLogger |
|
|
| torch.backends.cudnn.enabled = True |
| torch.backends.cudnn.benchmark = True |
|
|
|
|
| def hydra_params_to_dotdict(hparams): |
| def _to_dot_dict(cfg): |
| res = {} |
| for k, v in cfg.items(): |
| if isinstance(v, omegaconf.DictConfig): |
| res.update( |
| {k + "." + subk: subv for subk, subv in _to_dot_dict(v).items()} |
| ) |
| elif isinstance(v, (str, int, float, bool)): |
| res[k] = v |
|
|
| return res |
|
|
| return _to_dot_dict(hparams) |
|
|
|
|
| @hydra.main("config/config.yaml") |
| def main(cfg): |
| model = hydra.utils.instantiate(cfg.task_model, hydra_params_to_dotdict(cfg)) |
|
|
| early_stop_callback = pl.callbacks.EarlyStopping(patience=5) |
| checkpoint_callback = pl.callbacks.ModelCheckpoint( |
| monitor="val_acc", |
| mode="max", |
| save_top_k=2, |
| filepath=os.path.join( |
| cfg.task_model.name, "{epoch}-{val_loss:.2f}-{val_acc:.3f}" |
| ), |
| verbose=True, |
| ) |
| trainer = pl.Trainer( |
| gpus=list(cfg.gpus), |
| max_epochs=cfg.epochs, |
| early_stop_callback=early_stop_callback, |
| checkpoint_callback=checkpoint_callback, |
| distributed_backend=cfg.distrib_backend, |
| ) |
|
|
| trainer.fit(model) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
