run.py

import os
import yaml
import argparse
import numpy as np
from pathlib import Path
from models import *
from experiment import VAEXperiment
import torch.backends.cudnn as cudnn
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning import seed_everything
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
from dataset import VAEDataset, RoveVAEDataset, BIOSCANVAEDataset, BIOSCANTreeVAEDataset, BIOSCANPMLDataset, SimDataModule
from pytorch_lightning.strategies import DDPStrategy


parser = argparse.ArgumentParser(description="Generic runner for VAE models")
parser.add_argument(
    "--config",
    "-c",
    dest="filename",
    metavar="FILE",
    help="path to the config file",
    default="configs/vae.yaml",
)

args = parser.parse_args()
with open(args.filename, "r") as file:
    try:
        config = yaml.safe_load(file)
    except yaml.YAMLError as exc:
        print(exc)


tb_logger = TensorBoardLogger(
    save_dir=config["logging_params"]["save_dir"],
    name=config["model_params"]["name"],
)

# For reproducibility
seed_everything(config["exp_params"]["manual_seed"], True)

model = vae_models[config["model_params"]["name"]](**config["model_params"])

if config["model_params"].get("checkpoint", None) is not None:
    experiment = VAEXperiment.load_from_checkpoint(
        config["model_params"]["checkpoint"],
        vae_model=model,
        params=config["exp_params"],
    )
else:
    experiment = VAEXperiment(model, config["exp_params"])

data = SimDataModule(
    **config["data_params"], pin_memory=config["trainer_params"]["accelerator"] == "gpu"
)

data.setup()
runner = Trainer(
    logger=tb_logger,
    callbacks=[
        LearningRateMonitor(),
        ModelCheckpoint(
            save_top_k=2,
            dirpath=os.path.join(tb_logger.log_dir, "checkpoints"),
            monitor="val_loss",
            save_last=True,
        ),
    ],
    strategy="ddp_find_unused_parameters_true",
    use_distributed_sampler=False,
    **config["trainer_params"],
)


Path(f"{tb_logger.log_dir}/Inputs").mkdir(exist_ok=True, parents=True)
Path(f"{tb_logger.log_dir}/Samples").mkdir(exist_ok=True, parents=True)
Path(f"{tb_logger.log_dir}/Reconstructions").mkdir(exist_ok=True, parents=True)
Path(f"{tb_logger.log_dir}/Latent").mkdir(exist_ok=True, parents=True)

config_fp = f"{tb_logger.log_dir}/config.yaml"
with open(config_fp, 'w') as outfile:
    yaml.dump(config, outfile, default_flow_style=False)

print(f"======= Training {config['model_params']['name']} =======")
runner.fit(experiment, datamodule=data)