run.py

import os
import yaml
import argparse
import numpy as np
from pathlib import Path
from models import *  # Assuming vae_models are defined here
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
# Update import: use DDPStrategy (PL 2.0) instead of the old DDPPlugin
from pytorch_lightning.strategies import DDPStrategy
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')
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'])
experiment = VAEXperiment(model, config['exp_params'])
data = VAEDataset(**config["data_params"], pin_memory=len(config['trainer_params'].get('gpus', [])) != 0)
data.setup()

# (Optional) If your config still uses "max_nb_epochs", update it to "max_epochs" for PL 2.0:
trainer_params = config['trainer_params']
if 'max_nb_epochs' in trainer_params:
    trainer_params['max_epochs'] = trainer_params.pop('max_nb_epochs')

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
        ),
    ],
    # Use the new DDP strategy with the desired parameter.
    strategy=DDPStrategy(find_unused_parameters=False),
    **trainer_params
)

# Ensure directories exist for saving samples and reconstructions
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)

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