| from pytorch_lightning import Trainer |
| from models import vae_models |
| from config import config |
| from pytorch_lightning.callbacks import LearningRateMonitor |
| from pytorch_lightning.loggers import TensorBoardLogger |
| import os |
| os.environ['KMP_DUPLICATE_LIB_OK']='True' |
|
|
|
|
| def make_model(config): |
| model_type = config.model_type |
| model_config = config.model_config |
|
|
| if model_type not in vae_models.keys(): |
| raise NotImplementedError("Model Architecture not implemented") |
| else: |
| return vae_models[model_type](**model_config.dict()) |
|
|
|
|
| if __name__ == "__main__": |
| model = make_model(config) |
| train_config = config.train_config |
| logger = TensorBoardLogger(**config.log_config.dict()) |
| trainer = Trainer(**train_config.dict(), logger=logger, |
| callbacks=LearningRateMonitor()) |
| if train_config.auto_lr_find: |
| lr_finder = trainer.tuner.lr_find(model) |
| new_lr = lr_finder.suggestion() |
| print("Learning Rate Chosen:", new_lr) |
| model.lr = new_lr |
| trainer.fit(model) |
| else: |
| trainer.fit(model) |
| if not os.path.isdir("./saved_models"): |
| os.mkdir("./saved_models") |
| trainer.save_checkpoint( |
| f"saved_models/{config.model_type}_alpha_{config.model_config.alpha}_dim_{config.model_config.hidden_size}.ckpt") |
|
|
