eval_vae.py

import argparse, datetime, yaml
import torch

from latent_verse.models.vqvae import AugVAE
from latent_verse.loader import ImageDataModule
from latent_verse.callbacks import ReconstructedImageLogger

import pytorch_lightning as pl
from pytorch_lightning import seed_everything
from pytorch_lightning import Trainer

if __name__ == "__main__":

    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")

    config_parser = parser = argparse.ArgumentParser(description='Training Config', add_help=False)
    parser.add_argument('-c', '--config', default='', type=str, metavar='FILE',
                    help='YAML config file specifying default arguments')


    parser = argparse.ArgumentParser(description='AugVAE Evaluation')

    #path configuration
    parser.add_argument('--train_dir', type=str, default='dataset/train/',
                    help='path to train dataset')
    parser.add_argument('--val_dir', type=str, default='dataset/val/',
                    help='path to val dataset')
    parser.add_argument('--test_dir', type=str, default='dataset/val/',
                    help='path to test dataset')                    
    parser.add_argument('--log_dir', type=str, default='results/',
                    help='path to save logs')
    parser.add_argument('--backup_dir', type=str, default='backups/',
                    help='path to save backups for sudden crash')                    
    parser.add_argument('--ckpt_path', type=str,
                    help='path to previous checkpoint') 
    parser.add_argument('--pretrained_path', type=str,
                    help='path to pretrained codebook') 


    #training configuration     
    parser.add_argument('--finetune', action='store_true', default=False,
                    help='finetune pretrained model')                                                       
    parser.add_argument('--backup', action='store_true', default=False,
                    help='save backup and load from backup if restart happens')      
    parser.add_argument('--backup_steps', type =int, default = 1000,
                    help='saves backup every n training steps') 
    parser.add_argument('--log_images', action='store_true', default=False,
                    help='log image outputs. not recommended for tpus')   
    parser.add_argument('--image_log_steps', type=int, default=1000,
                    help='log image outputs for every n step. not recommended for tpus')   
    parser.add_argument('--refresh_rate', type=int, default=1,
                    help='progress bar refresh rate')    
    parser.add_argument('--precision', type=int, default=32,
                    help='precision for training')         

    parser.add_argument('--fake_data', action='store_true', default=False,
                    help='using fake_data for debugging') 
                                             
             
    parser.add_argument('--seed', type=int, default=42,
                    help='random seed')  
    parser.add_argument('--gpus', type=int, default=1,
                    help='number of gpus')     
    parser.add_argument('--gpu_dist', action='store_true', default=False,
                    help='distributed training with gpus') 

                                
    parser.add_argument('--num_sanity_val_steps', type=int, default=0,
                    help='num_sanity_val_steps')  
    parser.add_argument('--val_percent_check', type=int, default=100,
                    help='num_val_percent')
    parser.add_argument('--learning_rate', default=4.5e-6, type=float,
                    help='base learning rate')
    parser.add_argument('--lr_decay', action='store_true', default=False,
                    help = 'use learning rate decay')

    parser.add_argument('--batch_size', type=int, default=8,
                    help='training settings')  
    parser.add_argument('--epochs', type=int, default=1,
                    help='training settings')                                    
    parser.add_argument('--num_workers', type=int, default=16,
                    help='training settings')   
    parser.add_argument('--img_size', type=int, default=256,
                    help='training settings')              
    parser.add_argument('--resize_ratio', type=float, default=0.75,
                    help='Random resized crop lower ratio')
                  
    parser.add_argument('--debug', action='store_true', default=False,
                    help='debug run')  
    parser.add_argument('--web_dataset',action='store_true', default=False,
                    help='enable web_dataset')  
    parser.add_argument('--dataset_size', nargs='+', type=int, default=[1e9],
                    help='training settings')

    #model configuration
    parser.add_argument('--model', type=str, default='vqvae')
    parser.add_argument('--use_attn', type=bool, default=False, help='use attention')
    parser.add_argument('--codebook_dim', type=int, default=256,
                    help='number of embedding dimension for codebook')       
    parser.add_argument('--num_tokens', type=int, default=1024,
                    help='codebook size')        
    parser.add_argument('--double_z', type=bool, default=False,
                    help='double z for encoder')
    parser.add_argument('--z_channels', type=int, default=256,
                    help='image latent feature dimension')
    parser.add_argument('--resolution', type=int, default=256,
                    help='image resolution')
    parser.add_argument('--in_channels', type=int, default=3,
                    help='input image channel')
    parser.add_argument('--out_channels', type=int, default=3,
                    help='output image channel')    
    parser.add_argument('--hidden_dim', type=int, default=128,
                    help='hidden dimension init size')  
    parser.add_argument('--ch_mult', nargs='+', type=int, default=[1,1,2,2,4],
                    help='resnet channel multiplier')  
    parser.add_argument('--num_res_blocks', type=int, default=2,
                    help='number of resnet blocks')                     
    parser.add_argument('--attn_resolutions', nargs='+', type=int, default=[16],
                    help='model settings')  
    parser.add_argument('--dropout', type=float, default=0.0,
                    help='model settings') 
    parser.add_argument('--quant_beta', type=float, default=0.25,
                    help='quantizer beta')                     
    parser.add_argument('--quant_ema_decay', type=float, default=0.99,
                    help='quantizer ema decay')
    parser.add_argument('--quant_ema_eps', type=float, default=1e-5,
                    help='quantizer ema epsilon')  
                              
    #loss configuration
    parser.add_argument('--loss_type', type=str, default='mse')
    parser.add_argument('--p_loss_weight', type = float, default=0.1,
                    help = 'Perceptual loss weight')                    
    parser.add_argument('--codebook_weight', type=float, default=1.0,
                    help='lossconfig') 

    
    args_config, remaining = config_parser.parse_known_args()
    if args_config.config:
        with open(args_config.config, 'r') as f:
            cfg = yaml.safe_load(f)
            parser.set_defaults(**cfg)

    # The main arg parser parses the rest of the args, the usual
    # defaults will have been overridden if config file specified.
    args = parser.parse_args(remaining)

    #Map dataset directory to test_dir
    args.train_dir = args.test_dir
    args.val_dir = args.test_dir

    #random seed fix
    seed_everything(args.seed)   

    tpus = None
    gpus = args.gpus      
    if args.gpu_dist:
        torch.distributed.init_process_group(backend='nccl') 
        args.world_size = torch.distributed.get_world_size()
    else:
        args.world_size = args.gpus    

    args.base_lr = args.learning_rate             
    args.learning_rate = args.learning_rate * args.world_size * args.batch_size
    

    datamodule = ImageDataModule(args.train_dir, args.val_dir, 
                            args.batch_size, args.num_workers, 
                            args.img_size, args.resize_ratio, 
                            args.fake_data, args.web_dataset,
                            world_size = args.world_size,
                            dataset_size = args.dataset_size)
    

    if args.finetune:
        model = AugVAE.load_from_checkpoint(args.ckpt_path, finetuned=True, 
                                            ft_attn_resolutions=args.attn_resolutions, 
                                            ft_loss_type = args.loss_type, 
                                            ft_args = args)
    else:
        model = AugVAE.load_from_checkpoint(args.ckpt_path)      
               
  
    del model.loss
    model.args.log_dir = args.log_dir

    default_root_dir = args.log_dir

    if args.debug:
        limit_train_batches = 100
        limit_test_batches = 100
        args.backup_steps = 10
        args.image_log_steps = 10
    else:
        limit_train_batches = 1.0
        limit_test_batches = 1.0   


    logger = pl.loggers.tensorboard.TensorBoardLogger(args.log_dir, name='vqvae')     


    trainer = Trainer(tpu_cores=tpus, gpus= gpus, default_root_dir=default_root_dir,
                        max_epochs=args.epochs, progress_bar_refresh_rate=args.refresh_rate,precision=args.precision,
                        accelerator='ddp', benchmark=True,                         
                        num_sanity_val_steps=args.num_sanity_val_steps,
                        limit_val_batches = args.val_percent_check,                     
                        limit_train_batches=limit_train_batches,limit_test_batches=limit_test_batches,                                                  
                        logger = logger)
                                 
    if args.log_images:
        trainer.callbacks.append(ReconstructedImageLogger(every_n_steps=args.image_log_steps, nrow=args.batch_size))  
        
    print("Setting batch size: {}".format(model.hparams.batch_size))
        

    trainer.test(model, datamodule=datamodule)