run.py

import os
import yaml
import argparse
import logging
from pathlib import Path
from models import *
from experiment import VAEXperiment
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import CSVLogger
try:
    from pytorch_lightning.utilities.seed import seed_everything
except Exception:
    from lightning_fabric.utilities.seed import seed_everything
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
from dataset import VAEDataset


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')
parser.add_argument('--inference',
                    action='store_true',
                    help='Run inference visualizations after training')
parser.add_argument('--tsne-dim',
                    type=int,
                    choices=[2, 3],
                    default=2,
                    help='t-SNE plot dimensionality (2 or 3)')
parser.add_argument('--tsne-max-samples',
                    type=int,
                    default=2000,
                    help='Max samples to use for t-SNE')
parser.add_argument('--class-samples',
                    type=int,
                    default=8,
                    help='Samples per class for reconstructions')
parser.add_argument('--k',
                    type=int,
                    default=5,
                    help='K samples per class for stochastic reparameterization')
parser.add_argument('--interp-steps',
                    type=int,
                    default=10,
                    help='Interpolation steps between two images of same class')

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


csv_logger =  CSVLogger(save_dir=config['logging_params']['save_dir'],
                        name=config['model_params']['name'],)

log_dir = csv_logger.log_dir
Path(log_dir).mkdir(parents=True, exist_ok=True)

py_logger = logging.getLogger("vae-runner")
py_logger.setLevel(logging.INFO)
py_logger.handlers = []
py_logger.propagate = False
formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] %(message)s")
file_handler = logging.FileHandler(os.path.join(log_dir, "run.log"))
file_handler.setFormatter(formatter)
stream_handler = logging.StreamHandler()
stream_handler.setFormatter(formatter)
py_logger.addHandler(file_handler)
py_logger.addHandler(stream_handler)

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

model_params = dict(config['model_params'])
if 'img_size' not in model_params:
    patch_size = config.get('data_params', {}).get('patch_size')
    if isinstance(patch_size, (list, tuple)):
        if patch_size:
            model_params['img_size'] = patch_size[0]
    elif patch_size is not None:
        model_params['img_size'] = patch_size

model = vae_models[model_params['name']](**model_params)
experiment = VAEXperiment(model,
                          config['exp_params'])

trainer_params = dict(config['trainer_params'])
gpus = trainer_params.pop('gpus', None)
num_devices = None
if gpus is not None and 'devices' not in trainer_params:
    if isinstance(gpus, (list, tuple)):
        num_devices = len(gpus)
        trainer_params['devices'] = gpus
    else:
        num_devices = int(gpus)
        trainer_params['devices'] = gpus
    if num_devices and num_devices > 0:
        trainer_params['accelerator'] = 'gpu'

if 'devices' in trainer_params and num_devices is None:
    devices = trainer_params['devices']
    if isinstance(devices, (list, tuple)):
        num_devices = len(devices)
    else:
        try:
            num_devices = int(devices)
        except (TypeError, ValueError):
            num_devices = None

if num_devices is not None and num_devices > 1:
    trainer_params.setdefault('strategy', 'ddp')

data = VAEDataset(**config["data_params"], pin_memory=(num_devices or 0) != 0)

data.setup()
if getattr(experiment, "automatic_optimization", True) is False:
    if "gradient_clip_val" in trainer_params:
        experiment.gradient_clip_val = trainer_params.get("gradient_clip_val")
        experiment.gradient_clip_algorithm = trainer_params.get("gradient_clip_algorithm", "norm")
        py_logger.warning("Moving gradient clipping into manual optimization.")
        trainer_params.pop("gradient_clip_val", None)
        trainer_params.pop("gradient_clip_algorithm", None)
runner = Trainer(logger=csv_logger,
                 callbacks=[
                     LearningRateMonitor(),
                     ModelCheckpoint(save_top_k=2,
                                     dirpath=os.path.join(log_dir, "checkpoints"),
                                     monitor="val_loss",
                                     save_last=True),
                 ],
                 **trainer_params)


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


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

if args.inference:
    py_logger.info("======= Running inference visualizations =======")
    experiment.run_inference(
        datamodule=data,
        output_dir=os.path.join(log_dir, "Inference"),
        tsne_dim=args.tsne_dim,
        tsne_max_samples=args.tsne_max_samples,
        class_samples=args.class_samples,
        reparam_k=args.k,
        interp_steps=args.interp_steps,
    )