eval_MoCoDAD.py

import argparse
import os

import pytorch_lightning as pl
import yaml
from models.mocodad import MoCoDAD
from models.mocodad_latent import MoCoDADlatent
try:
    from pytorch_lightning.loggers import WandbLogger
except Exception:
    WandbLogger = None
from utils.argparser import init_args
from utils.dataset import get_dataset_and_loader


def build_logger(args: argparse.Namespace):
    if not getattr(args, "use_wandb", False):
        return False
    if WandbLogger is None:
        raise ImportError("use_wandb=true but WandbLogger is unavailable. Install wandb and pytorch_lightning loggers.")
    return WandbLogger(
        project=args.project_name,
        group=args.group_name,
        entity=args.wandb_entity,
        name=f"{args.dir_name}_{args.split}_eval",
        config=vars(args),
    )


def resolve_eval_devices(devices):
    if isinstance(devices, (list, tuple)):
        return devices[:1]
    return devices


if __name__== '__main__':
    
    # Parse command line arguments and load config file
    parser = argparse.ArgumentParser(description='MoCoDAD')
    parser.add_argument('-c', '--config', type=str, required=True)
    args = parser.parse_args()
    args = yaml.load(open(args.config), Loader=yaml.FullLoader)
    args = argparse.Namespace(**args)
    args = init_args(args)

    # Initialize the model
    model = MoCoDADlatent(args) if hasattr(args, 'diffusion_on_latent') else MoCoDAD(args)
    wandb_logger = build_logger(args)
    
    if args.load_tensors:
        # Load tensors and test
        auc = model.test_on_saved_tensors(split_name=args.split)
        if wandb_logger:
            wandb_logger.experiment.log({"AUC": auc})
            wandb_logger.experiment.finish()
    else:
        # Load test data
        print('Loading data and creating loaders.....')
        ckpt_path = os.path.join(args.ckpt_dir, args.load_ckpt)
        dataset, loader, _, _ = get_dataset_and_loader(args, split=args.split)
        
        # Initialize trainer and test
        trainer = pl.Trainer(accelerator=args.accelerator, devices=resolve_eval_devices(args.devices),
                             default_root_dir=args.ckpt_dir, max_epochs=1, logger=wandb_logger)
        try:
            out = trainer.test(model, dataloaders=loader, ckpt_path=ckpt_path, weights_only=False)
        except TypeError:
            out = trainer.test(model, dataloaders=loader, ckpt_path=ckpt_path)